mlir::dyn_cast<RankedTensorType>(operand_type)) {
      tensorflow::TensorShapeProto shape_proto;
      ConvertToTensorShapeProto(ranked_tensor_type.getShape(), &shape_proto);
      *arg->mutable_shape() = std::move(shape_proto);
    } else {
      arg->mutable_shape()->set_unknown_rank(true);
    }

    if (failed(SetOpSharding(op, input_shardings.getValue()[index],

             std::vector<XlaShapeLayoutHelpers::ShapeDeterminationFns>
                 shape_determination_fns)
        : jit_device_type_(jit_device_type),
          shape_determination_fns_(std::move(shape_determination_fns)) {}

    // The index of the device on this host.
    int device_ordinal() const;

    const DeviceType& jit_device_type() const { return jit_device_type_; }

      if (!mlir::cast<BoolAttr>(legalize_tfl_variables_attr).getValue()) return;
    }

    RewritePatternSet patterns(&getContext());
    populateWithGenerated(patterns);
    (void)applyPatternsAndFoldGreedily(module, std::move(patterns));
  }
};

}  // namespace

std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeVariablesPass() {
  return std::make_unique<LegalizeVariablesPass>();
}

}  // namespace TFL

    auto signedRhs = addSign(std::get<1>(zip), etype);
    FailureOr<decltype(signedLhs)> result(Convert()(signedLhs, signedRhs));
    if (failed(result)) {
      return {};
    }
    values.push_back(std::move(*result));
  }

  return DenseElementsAttr::get(type, values);
}

template <typename BinaryOpType>
class FoldBroadcastInDimBeforeBinaryElementwiseOp
    : public OpRewritePattern<BinaryOpType> {

                                    // ignore
                                }
                                Files.move(path, newPath);
                                logger.info("Move {} to {}", path, newPath);
                            } catch (final IOException e) {
                                logger.warn("Failed to move {}", path, e);
                            }
                        }
                    }
                });

	ADD $-1, R11, R12
	ADD $-1, INP
	ADD $-1, OUT
	PCALIGN $16
looptail_vsx:
	// Copying the result to OUT
	// in bytes.
	MOVBZU 1(R12), KEY
	MOVBZU 1(INP), TMP
	XOR    KEY, TMP, KEY
	MOVBU  KEY, 1(OUT)
	BDNZ   looptail_vsx

	// Clear the stack values
	STXVW4X VS48, (R11)(R0)
	STXVW4X VS48, (R11)(R8)
	STXVW4X VS48, (R11)(R9)
	STXVW4X VS48, (R11)(R10)

	CMP	R5, R6;				\
	BNE	-3(PC);				\
	AND	$0xf, R4, R6;			\
	CBZ	R6, 6(PC);			\
	/* if R6=(argsize&15) != 0 */		\
	ADD	R6, R5, R6;			\
	/* copy 1 byte a time for the rest */	\
	MOVBU.P	1(R3), R7;			\
	MOVBU.P	R7, 1(R5);			\
	CMP	R5, R6;				\
	BNE	-3(PC);				\
	/* set up argument registers */		\
	MOVD	regArgs+40(FP), R20;		\
	CALL	·unspillArgs(SB);		\
	/* call function */			\
	MOVD	f+8(FP), R26;			\

  let summary = "Sets the program key for embedding ops.";
  let constructor = "TFDevice::CreateEmbeddingProgramKeyPass()";
    let description = [{
    Passes in the program key to embedding ops. Will move the embedding ops
    after a _TPUCompileMlir op if there is no predecessor _TPUCompileMlir op.
    Both the embedding op and compile op are assumed to be wrapped in separate

  patterns.add<ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulOp>,
               ConvertTFBatchMatMulToEinsumOp<TF::BatchMatMulV2Op>>(
      &getContext());
  (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
}

}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateBatchMatMulToEinsumPass() {
  return std::make_unique<BatchMatMulToEinsumPass>();
}

}  // namespace TF

	MOVV	REGCTXT, R4
	MOVV	R25, R5
	JAL	·moveMakeFuncArgPtrs<ABIInternal>(SB)
	MOVV	32(R3), REGCTXT // restore REGCTXT
	MOVV	REGCTXT, 8(R3)
	MOVV	$argframe+0(FP), R20
	MOVV	R20, 16(R3)
	MOVB	R0, LOCAL_RETVALID(R3)
	ADDV	$LOCAL_RETVALID, R3, R20
	MOVV	R20, 24(R3)
	ADDV	$LOCAL_REGARGS, R3, R20
	MOVV	R20, 32(R3) // frame size to 32+SP as callreflect args)
	JAL	·callMethod(SB)