builder.getBoolAttr(true));
  ConstOp shape = builder.create<ConstOp>(loc, shape_attr);
  shape->setAttr(kICIWeightDistributionMlirBridgeMarker,
                 builder.getBoolAttr(true));
  FillOp fill = builder.create<FillOp>(loc, shape, zero);
  fill->setAttr(kICIWeightDistributionMlirBridgeMarker,
                builder.getBoolAttr(true));
  return fill;
}

    Type expected = fn_type.getInput(i);
    if (val.getType() != expected) {
      val =
          builder->create<TF::CastOp>(loc, expected, val,
                                      /*Truncate=*/builder->getBoolAttr(false));
    }
    operands.push_back(val);
  }
  return builder->create<func::CallOp>(loc, fn, operands).getOperation();
}

// Prepares for jump to the given block by introducing necessary tensor_cast

                                               rewriter.getF32Type());
    ::mlir::Value mean = rewriter.create<TF::MeanOp>(
        bn_op.getLoc(), mean_var_type, inputs, reduce_dim_op,
        /*keep_dims=*/rewriter.getBoolAttr(false));

    // Compute variance
    Value shape_value =
        getShapeValue(bn_op.getLoc(), bn_op.getOperand(), rewriter);
    auto broadcast_mean = broadcastToFeatureDim(

    return rewriter.createOrFold<TF::CastOp>(loc, new_type, val,
                                             rewriter.getBoolAttr(false));
  }
  return rewriter.createOrFold<TF::CastOp>(
      loc, UnrankedTensorType::get(new_ele_type), val,
      rewriter.getBoolAttr(false));
}

// Utility function to-
// 1. Create a tfl.const op with an int32_t values, from an MLIR Value, if the

  // in TensorFlow and allows lowering to V1 control flow for loop
  // parallelization.
  dataset_while->setAttr("_lower_using_switch_merge",
                         builder.getBoolAttr(true));

  return dataset_while;
}

// Populate the cond of `dataset_while`.  The cond body just returns the
// condition of whether to continue to next iteration.

  auto matmul_type =
      RankedTensorType::get(matmul_shape, original_type.getElementType());
  Value out = rewriter.create<TF::BatchMatMulV2Op>(
      op.getLoc(), matmul_type, lhs, rhs, rewriter.getBoolAttr(false),
      rewriter.getBoolAttr(false));

  bool out_reshape_need = (reshape_shape.size() != matmul_shape.size() ||
                           original_type.getRank() != matmul_shape.size());

def convertIntAttrTo32Bit : NativeCodeCall<
    "$_builder.getI32IntegerAttr($0.cast<IntegerAttr>().getInt())">;

// Builds a constant bool attribute.
class GetBoolAttr<int value> :
    NativeCodeCall<"$_builder.getBoolAttr(" # value #")">;

// Converts an integer attribute $0 to 64-bit with builder.
def convertIntAttrTo64Bit : NativeCodeCall<
    "$_builder.getI64IntegerAttr($0.cast<IntegerAttr>().getInt())">;

      loc, tensor_type.clone(builder.getIntegerType(32)), tensor);
  auto reduced =
      builder.create<TF::SumOp>(loc, tensor_i32, reduction_indices_value,
                                /*keep_dims=*/builder.getBoolAttr(true));
  auto mul_op = builder.create<TF::MulOp>(loc, zp, reduced);

  SmallVector<Value> folded_results = ConstantFoldOpIfPossible(mul_op);
  return folded_results.front();
}

    Type expected_type = std::get<1>(ArgAndType);
    if (arg.getType() != expected_type) {
      arg = builder.create<CastOp>(loc, expected_type, arg,
                                   /*Truncate=*/builder.getBoolAttr(false));
    }
    casted_args.push_back(arg);
  }
  auto call = builder.create<func::CallOp>(loc, func, casted_args);

  auto results = call.getResults();
  auto block_args = entry->getArguments();

      if (operand.getType() != expected_type) {
        operand = rewriter.create<TF::CastOp>(
            op.getLoc(), expected_type, operand,
            /*Truncate=*/rewriter.getBoolAttr(false));
      }
      casted_operands.push_back(operand);
    }

    auto call = rewriter.create<func::CallOp>(
        op->getLoc(), main_fn.getSymName(), main_fn.getResultTypes(),
        casted_operands);