from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common


def mnist_model():
  """Creates a MNIST model."""
  model = tf.keras.models.Sequential()
  model.add(tf.keras.layers.Flatten())
  model.add(tf.keras.layers.Dense(128, activation='relu'))
  model.add(tf.keras.layers.Dense(10, activation='softmax'))
  return model


class TestModule(tf.Module):

  def __init__(self):

  "TFL::IsReducedTailOfShape($0.getType(), $1.getType())">>;

def IsRankLessThanEqualTo : Constraint<CPred<
  "$0.getType().cast<ShapedType>().getRank() <= "
  "$1.getType().cast<ShapedType>().getRank()">>;

def Flatten : NativeCodeCall<
  "$0.cast<DenseElementsAttr>()"
    ".reshape(RankedTensorType::get({$0.getType().cast<ShapedType>().getNumElements()}, "
                                   "$0.getType().cast<ShapedType>().getElementType()))">;

// Additionally, we take advantage of the fact that the SignatureDefFunction's
// associated functiondef has lexicographically ordered inputs/outputs due to
// nest.flatten.
Status LoadSignatureDefFunctionMetadata(
    const SavedConcreteFunction& saved_concrete_function,
    SignatureDefFunctionMetadata* out) {
  std::vector<SignatureDefParam> args;

  auto indices_size = ExpandDims(scope, Size(scope, indices), 0);
  Output outer_shape, flat_values_shape;
  if (batch_dims != 0) {
    auto values_shape = Shape(scope, values);
    // Add the batch offsets to indices and flatten the batch dimensions.
    outer_shape = Slice(scope, values_shape, {0}, {batch_dims});
    auto inner_shape =
        Slice(scope, Slice(scope, values_shape, {batch_dims}, {-1}), {1}, {-1});

      rhs,
      DenseIntElementsAttr::get(
          RankedTensorType::get({rhs_rank}, rewriter.getI64Type()),
          rhs_permutation));
  // Reshapes lhs to flatten out_dimensions and contracting_dimensions.
  llvm::SmallVector<int64_t, 4> lhs_flattened_shape = Concat<int64_t>(
      lhs_dot_dimensions_info.batch_dimensions().SizesArray(),
      llvm::ArrayRef<int64_t>{

// TODO(b/190203981): Move to a separate nest-like library.
void Flatten(const TaggedValue& value,
             std::vector<AbstractTensorHandle*>* flat_args) {
  if (value.type() == TaggedValue::Type::TENSOR) {
    flat_args->emplace_back(value.tensor().get());
  } else if (value.type() == TaggedValue::Type::TUPLE) {
    for (const auto& t : value.tuple()) {
      Flatten(t, flat_args);
    }
  } else {

    llvm::StringRef equation, RankedTensorType lhs_ty) {
  llvm::StringRef lhs;
  llvm::StringRef out;
  std::tie(lhs, out) = equation.split("->");
  if (lhs.empty() || out.empty()) return std::nullopt;

  // Try to flatten the "..." if possible.
  int lhs_named_label, rhs_named_label;

  // following rhs and rhs_ty variables are non-functional here only created to
  // comply with the existing API
  llvm::StringRef rhs;

  def f0004_list_2_elements(self):
    return [[tf.constant(1.0, shape=[1]), tf.constant(1.0, shape=[2])]]

  # Check index paths for dicts.
  # Keys are linearized in sorted order, matching `tf.nest.flatten`.
  # More thorough testing of this is in structured_input.py. The underlying code
  # path for linearization is shared, so no need to replicate that testing here.
  #
  # CHECK:      func {{@[a-zA-Z_0-9]+}}() -> (

         (std::equal(i1, reduced_e1, i2));
}

// Check if the value of the last dimension of type1 is equal to the number of
// elements in type2. This is a required condition to flatten type2 to form a
// 1D array and allow the binaryOp handle the broadcasting implicitly.
bool IsLastDimEqualToNumElements(Type type1, Type type2) {
  return (mlir::cast<ShapedType>(type1).getRank() >= 1 &&

    const SavedConcreteFunction& saved_concrete_function,
    const FunctionDef* function_def) {
  // tf.functions go through many transformations before becoming FunctionDefs
  // 1. flatten user-provided inputs:
  // https://github.com/tensorflow/tensorflow/blob/1c064ab76064c58e54261b805027474885a1534d/tensorflow/python/eager/function.py#L2671-L2675
  // 2. convert user-provided inputs to tensors: