ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());
    for (size_t i = 0; i < permutation.size(); ++i)
      new_shape[i] = shape[permutation[i]];

    return RankedTensorType::get(new_shape, ranked_type.getElementType());
  }

  return type;
}

    auto new_shape_const_attr =
        DenseElementsAttr::get(shape_spec_type, new_shape.getShape());
    rewriter.setInsertionPointAfter(weight_op);
    auto new_shape_const = rewriter.create<arith::ConstantOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(
        weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);

                                        llvm::ArrayRef<int32_t> perm) {
  llvm::SmallVector<int64_t> new_shape(shape.size());
  for (const auto &perm_element : enumerate(perm)) {
    new_shape[perm_element.index()] = shape[perm_element.value()];
  }
  return new_shape;
}

// Determine if op commutes with transposes. Requires a strict
// definition of Elementwise, all i/o shapes and types must be same-rank

      //   [H, W, ..., filter_in_depth, G, out_depth / G].
      auto new_shape = llvm::to_vector<6>(filter_shape);
      new_shape.back() = feature_group_count;
      new_shape.push_back(filter_shape.back() / feature_group_count);
      Type filter_element_ty = filter_ty.getElementType();
      auto ty =
          tensorflow::GetTypeFromTFTensorShape(new_shape, filter_element_ty);

          value_attr.getValues<Attribute>()[outer_index * inner_dims_size + i]);
    }
  }

  std::vector<int64_t> new_shape = value_attr.getType().getShape().vec();
  new_shape[0] = lookup_attr.getType().getShape()[0];
  auto new_type = value_attr.getType().clone(new_shape);

  return DenseElementsAttr::get(new_type, new_values);
}

    } else {
      output_ty = tensorflow::GetTypeFromTFTensorShape(
          {static_cast<int64_t>(sub_shape.size())}, output_elt_ty);
    }
  }

  // Down-cast to 32 bit int if needed.
  if (output_elt_ty.isInteger(32)) {
    SmallVector<int32_t, 2> sub_shape_i32(sub_shape.size());
    std::transform(sub_shape.begin(), sub_shape.end(), sub_shape_i32.begin(),
                   [](int64_t d) { return static_cast<int32_t>(d); });

      return {};
    }
  } else if ((new_shape.size() == shape.size() + 1) && new_shape.front() == 1) {
    // Handle the [A, B, C] -> [1, A, B, C] reshape case.
    if (!(std::equal(shape.begin(), shape.end(), new_shape.begin() + 1) &&
          quant_dim == new_shape.size() - 1)) {
      return {};
    }
  } else {
    return {};
  }

    %2 = "tf.PartitionedCall"(%1) {config = "", config_proto = "", executor_type = "", f = @get_shape} : (tensor<?x?xf32>) -> (tensor<?xi64>)

    // CHECK: return %[[RESULT]]
    func.return %2 : tensor<?xi64>
  }

  // CHECK-LABEL: func @get_shape
  func.func @get_shape(%arg0 : tensor<*xi64>) -> tensor<?xi64> {
    %0 = "tf.Shape"(%arg0) : (tensor<*xi64>) -> tensor<?xi64>

                                       /*shape=*/new_shape,
                                       /*type=*/new_filter_quantized_type);
    } else if (element_type.isa<UniformQuantizedType>()) {
      auto per_tensor_type = element_type.cast<UniformQuantizedType>();
      new_filter_result_type =
          RankedTensorType::getChecked(location,
                                       /*shape=*/new_shape,

    if (!input_ty.hasRank() || input_ty.getRank() != 4) {
      return failure();
    }

    int64_t batch_cst = input_ty.getShape()[0];
    int64_t channels_cst = input_ty.getShape()[3];

    int64_t in_y_cst = input_ty.getShape()[1];
    int64_t in_x_cst = input_ty.getShape()[2];
    int64_t in_spatial_cst =
        in_y_cst < 0 || in_x_cst < 0 ? -1 : in_y_cst * in_x_cst;