}

def TFL_LogicalNotOp : TFL_Op<"logical_not", [
    Pure,
    SameOperandsAndResultShape]> {
  let summary = "Logical NOT operator";

  let description = [{
    Element-wise logical NOT operation.
  }];

  let arguments = (ins TFL_BoolTensor:$lhs);

  let results = (outs TFL_BoolTensor:$output);
}

def TFL_LogicalOrOp : TFL_Op<"logical_or", [Pure]> {

            LogicOperationKind.OR -> KaContractBinaryLogicExpression.KaLogicOperation.OR
        }
    )

    override fun visitLogicalNot(logicalNot: ConeLogicalNot, data: Unit): KaContractLogicalNotExpression =
        KaContractLogicalNotExpression(logicalNot.arg.accept())

    override fun visitIsInstancePredicate(isInstancePredicate: ConeIsInstancePredicate, data: Unit): KaContractIsInstancePredicateExpression =

        %1 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
        %2 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 1} : () -> tensor<i64>
        tf_device.return
      }
      tf_executor.yield
    }
    tf_executor.fetch
  }
  func.return
}

                   TF::EnqueueTPUEmbeddingArbitraryTensorBatchOp>(op);
}

// Returns the device ordinal (`device_ordinal`) for a replica (`replica_id`)
// and logical core (`logical_core`).
// `replica_id` is the index of the ancestor ReplicateOp in [0, num_replicas).
// `logical_core` is the index of the TPU core in [0, num_cores_per_replica).
// `device_ordinal` is the index of the TPU core relative to its host.

  xla::DeviceAssignment device_assignment(num_replicas, num_cores_per_replica);
  int pos = 0;
  for (int replica = 0; replica < num_replicas; ++replica) {
    for (int logical_core = 0; logical_core < num_cores_per_replica;
         ++logical_core) {
      int x = device_assignment_attr[pos++];
      int y = device_assignment_attr[pos++];
      int z = device_assignment_attr[pos++];
      int core = device_assignment_attr[pos++];

      %1:2 = tf_device.replicate {n = 2 : i32, devices = {TPU_REPLICATED_CORE_0 = ["/job:worker/replica:0/task:0/device:TPU:1", "/job:worker/replica:0/task:0/device:TPU:2"]}} {
        %2 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
        tf_device.return %2 : tensor<i64>
      }
      tf_executor.yield %1#0, %1#1 : tensor<i64>, tensor<i64>
    }
    tf_executor.fetch
  }
  func.return
}

  %105 = "tf.LogicalOr"(%104, %2) {device = ""} : (tensor<i1>, tensor<i1>) -> tensor<i1>
  %106 = "tf.Equal"(%103, %13) {device = "", incompatible_shape_error = true} : (tensor<i64>, tensor<i64>) -> tensor<i1>
  %107 = "tf.LogicalOr"(%105, %106) {device = ""} : (tensor<i1>, tensor<i1>) -> tensor<i1>

REGISTER_NO_GRADIENT_OP("Equal");
REGISTER_NO_GRADIENT_OP("ApproximateEqual");
REGISTER_NO_GRADIENT_OP("NotEqual");
REGISTER_NO_GRADIENT_OP("LogicalAnd");
REGISTER_NO_GRADIENT_OP("LogicalOr");
REGISTER_NO_GRADIENT_OP("LogicalNot");
REGISTER_NO_GRADIENT_OP("Floor");

// Conjugate helper function returns the conjugate of an Output if it
// is complex valued.
Output ConjugateHelper(const Scope& scope, const Output& out) {

  // CHECK: tf_device.replicate
  // CHECK:   tf._TPUDeviceOrdinalPlaceholder
  %0:2 = tf_device.replicate([%arg0, %arg1] as %ri: tensor<*xf32>) {n = 2: i32} {
    %1 = "tf._TPUDeviceOrdinalPlaceholder"() {logical_core = 0} : () -> tensor<i64>
    tf_device.return %1 : tensor<i64>
  }
  func.return
}


// CHECK-LABEL: func @replicate_resource_var_arg_shape

    (TF_ConstOp ConstantAttr<RankedF32ElementsAttr<[]>, "1.0f">))),
  (TF_Log1pOp:$dest $arg), [], [(CopyAttrs $src, $dest)]>;

//===----------------------------------------------------------------------===//
// LogicalNot op patterns.
//===----------------------------------------------------------------------===//

def LogicalNotOfEqual : Pat<
  (TF_LogicalNotOp:$src (TF_EqualOp $arg0, $arg1, $shape_error)),