// expected-error@-1 {{'tf_executor.Merge' op expects all operands to be broadcastable with output type but got 'tensor<i1>' vs 'tensor<*xf32>'}}
    tf_executor.fetch %value : tensor<*xf32>
  }
  func.return %result : tensor<*xf32>
}

// -----

// Check that merge data inputs are broadcastable to the output
func.func @invalid_merge(%arg0: tensor<*xf32>, %arg1: tensor<4xf32>) -> tensor<8xf32> {

    // Scalar identity is broadcastable to any operand shape, we only need to
    // check that operand has the same shape as a result.
    bool scalar_identity = identity_ty.hasRank() && identity_ty.getRank() == 0;
    if (scalar_identity) return operand_ty == result_ty;

    // If identity is not a scalar, we must verify that identity shape is
    // statically known to be broadcastable to the operand shape and the operand

// Canonicalize: Log(1.0 + x) to Log1p(x)
//
// We currently do this rewrite only if the constant `1` is a scalar, because
// it is safely broadcastable to any shape. To be able to canonicalize when
// constant values is not a scalar, we have to first prove that it is
// broadcastable to `x`, which requires static shape information.
def LogToLog1p : Pat<
  (TF_LogOp:$src (TF_AddV2Op $arg,

}

def TfExecutor_SwitchOp : TfExecutor_Op<"Switch",
    [ControlOperandsAfterAllData, HasParent<"GraphOp">,
     PredOpTrait<"data operand must be broadcastable to true result",
                 TF_OpIsBroadcastableToRes<0, 0>>,
     PredOpTrait<"data operand must be broadcastable to false result",
                 TF_OpIsBroadcastableToRes<0, 1>>]>{
  let summary = [{

      return failure();
    }
  }
  // Ensure that batch shapes are broadcastable.
  tensorflow::MatMulBCast bcast(
      absl::InlinedVector<int64_t, 4>(lhs_shape.begin(), lhs_shape.end()),
      absl::InlinedVector<int64_t, 4>(rhs_shape.begin(), rhs_shape.end()));

  if (!bcast.IsValid()) {
    // Input batch dimensions must be broadcastable
    return failure();
  }

#include "mlir/IR/Operation.h"  // from @llvm-project
#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace TFL {

// For add/mul/div/sub and other broadcastable ops.
class ArithmeticCountUtilHelper {
 public:
  static bool GetFirstOutputCount(mlir::Operation* op, int64_t* count) {
    auto output = op->getResult(0);
    auto output_type =

// the core implementation requiring SameOperandsAndResultType.
//
// This shouldn't be used for side effecting ops.
def TF_Involution : NativeOpTrait<"TF::IsInvolution">;

// Variant of broadcastable trait that considers TF's subtype behavior.
class TF_OpIsBroadcastableToRes<int opId, int resId> : And<[
    TCOpResIsShapedTypePred<opId, resId>,
    CPred<"mlir::tf_type::BroadcastCompatible("

    }

    // Check if alpha is broadcastable
    for (int i = 0; i < alpha_type.getRank(); i++) {
      if (alpha_type.getDimSize(i) != input_type.getDimSize(i + 1) &&
          alpha_type.getDimSize(i) != 1) {
        return op.emitOpError(
            llvm::formatv("'alpha' is not broadcastable at dimension {0}.", i));
      }
    }
  }

        if (shape_x.size() < 2 || shape_y.size() < 2) {
          return false;
        }

        // Checks outer dimensions (i.e., the dimensions higher than 2D) are
        // broadcastable. If true, then get the broadcasted shape for outer
        // dimension.
        if (!OpTrait::util::getBroadcastedShape(
                shape_x.drop_back(2), shape_y.drop_back(2), result_shape)) {
          return false;

  });
}

bool IsBroadcastableElementsAttrs(mlir::TypedAttr a, mlir::TypedAttr b) {
  // This would return false if we had unranked tensors (where they should
  // probably be considered as broadcastable), but given we are working with
  // attributes here that shouldn't be an issue,
  return OpTrait::util::getBroadcastedType(a.getType(), b.getType()) != Type();
}