// Relu op patterns.
//===----------------------------------------------------------------------===//

// TODO(hinsu): Make these patterns to TF to TF lowering. Relu6 lowering will
// require HLO canonicalization of min and max on a tensor to ClampOp.

// TODO(hinsu): Lower quantized types after supporting them in GetScalarOfType.
def : Pat<(TF_ReluOp AnyTensor:$input),
          (CHLO_BroadcastMaxOp

    if (auto sharding =
            func.getArgAttrOfType<mlir::StringAttr>(i, kShardingAttr)) {
      if (!sharding.getValue().empty()) {
        BlockArgument arg = func.getArgument(i);
        // TODO(hinsu): Instead of setting both 'sharding' and '_XlaSharding'
        // attributes, only set the 'sharding' attribute. Both attributes are
        // currently required as the XlaSharding xla op kernel doesn't use the

  // to empty elements attribute. This is restricted to the numerical element
  // types as the DenseElementsAttr only supports numerical and string types.
  // TODO(hinsu): Handle ops that have one of the results empty for constant
  // propagation.
  bool has_empty_numerical_results =
      llvm::all_of(inst->getResultTypes(), [](Type ty) {

/// list.
void PopulateLegalizeTfPatterns(MLIRContext* context,
                                RewritePatternSet* patterns);

// Populates TF to MHLO legalization for some of the quantization ops.
//
// TODO(hinsu): Remove this once we combine quantized and non quantized op
// legalization in the ODML conversion pipeline.
void PopulateLegalizeTfQuantizationPatterns(MLIRContext* context,

//
// Computation of the reduction axis for the Sum op depends on whether the
// input is a scalar or not. Restrict pattern to ranked inputs so that input to
// the Sum op is also ranked.

// TODO(hinsu): Support scalar inputs by introducing reshape to 1D.
def NonScalarType : Type<Neg<HasAnyRankOfPred<[0]>>, "Non scalar type">;

def LowerSoftmaxCrossEntropyWithLogitsOp : Pattern<

      return success();
    }

    // Extract the constant cond value.
    DenseElementsAttr cond;
    if (!matchPattern(op.getCond(), m_Constant(&cond))) return failure();

    // TODO(hinsu): Handle constants that are not scalar booleans.
    auto cond_type = mlir::dyn_cast<RankedTensorType>(cond.getType());
    if (!cond_type || !cond_type.getShape().equals({}) ||

  return *ops;
}

bool IsOpTypeAllowedTf2XlaFallback(const TypeID& type_id) {
  // Allowlisted TensorFlow ops are known to have well behaved tf2xla kernels
  // building valid MLIR using MlirHloBuilder.
  // TODO(hinsu): Drop explicit allowlist when MLIR based bridge is enabled for
  // all tf2xla kernels.
  // Use a pointer for the static set, so the set is not destructed upon thread
  // end, which would not be thread safe.

  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  // The SCCP pass performs constant propagation across the IR, which, for
  // example, propagates constant arguments into callee functions.
  // TODO(hinsu): Investigate if we really need SCCP pass before shape inference
  // and can do with just one pass after the shape inference.
  pm.addPass(mlir::createSCCPPass());

    case DT_FLOAT8_E4M3FN:
      return ConvertTensorOfCustomFloatType(input_tensor, type);
    case DT_STRING:
      return ConvertStringTensor(input_tensor, type);
    default:
      // TODO(hinsu): Remove mangling now that there is a special attribute.
      return ElementsAttr(
          mlir::TF::TensorProtoAttr::get(type, MangleTensor(input_tensor)));
  }

#undef CONVERT_FLAT
}

      attr.setValue(ConvertToDenseElementsAttr(
          mlir::cast<ArrayAttr>(attr.getValue()), rewriter));
      converted_attrs.push_back(attr);
    }
  }
  return converted_attrs;
}

// TODO(hinsu): Move this pattern to legalize_tf after resolving the dependency
// on the tensor proto.
class ConvertUniformQuantizedDotHybridOp
    : public OpConversionPattern<TF::UniformQuantizedDotHybridOp> {
 public: