StrAttr:$mini_batch_in_csr
  );

  let results = (outs
    TF_Int32Tensor:$row_pointers,
    TF_Int32Tensor:$sorted_sample_ids,
    TF_Int32Tensor:$sorted_token_ids,
    TF_Float32Tensor:$sorted_gains,
    TF_Int32Tensor:$row_pointers_unpadded_size,
    TF_Int32Tensor:$ids_unpadded_size,
    TF_Int32Tensor:$num_minibatches_per_physical_sparse_core
  );
}

  let arguments = (ins
    TensorOf<[TF_Float32, TF_Qint8]>:$lhs,
    TensorOf<[TF_Qint8]>:$rhs,
    TF_Float32Tensor:$lhs_scales,
    TF_Int32Tensor:$lhs_zps,
    TF_Float32Tensor:$rhs_scales,
    TF_Int32Tensor:$rhs_zps,
    TF_Float32Tensor:$output_scales,
    TF_Int32Tensor:$output_zps,

    StrAttr:$dimension_numbers,
    DefaultValuedOptionalAttr<I64Attr, "-1">:$lhs_quantization_axis,

  let summary = "Assign group keys based on group assignment.";

  let arguments = (ins
    TF_Int32Tensor:$group_assignment,
    TF_Int32Tensor:$device_index,
    TF_Int32Tensor:$base_key
  );

  let results = (outs
    TF_Int32Tensor:$group_size,
    TF_Int32Tensor:$group_key
  );
}

def TF_CollectiveBcastRecvOp : TF_Op<"CollectiveBcastRecv", []> {

def : Pat<(MHLO_NotOp TF_Int16Tensor:$input),
            (TF_BitwiseXorOp $input,
              (TF_ConstOp (GetRankedScalarAttr<"", 16, "", "-1">)))>;
def : Pat<(MHLO_NotOp TF_Int32Tensor:$input),
            (TF_BitwiseXorOp $input,
              (TF_ConstOp (GetRankedScalarAttr<"", 32, "", "-1">)))>;
def : Pat<(MHLO_NotOp TF_Int64Tensor:$input),
            (TF_BitwiseXorOp $input,

def TF_BoolTensor : TensorOf<[TF_Bool]>;

def TF_IntTensor : TensorOf<[TF_Int]>;
def TF_Int4Tensor : TensorOf<[TF_Int4]>;
def TF_Int8Tensor : TensorOf<[TF_Int8]>;
def TF_Int16Tensor : TensorOf<[TF_Int16]>;
def TF_Int32Tensor : TensorOf<[TF_Int32]>;
def TF_Int64Tensor : TensorOf<[TF_Int64]>;
def TF_I32OrI64Tensor : TensorOf<[TF_I32OrI64]>;

def TF_Uint4Tensor : TensorOf<[TF_Uint4]>;
def TF_Uint8Tensor : TensorOf<[TF_Uint8]>;