added SizeNe

Author: miladm

torch_xla/csrc/ops/dynamic_ir.cpp

@@ -112,6 +112,25 @@ int64_t SizeEq::getDynamicValue() const {
 
 std::string SizeEq::ToString() const { return "SizeEq"; }
 
+SizeNe::SizeNe(torch::lazy::Value a, torch::lazy::Value b)
+    : XlaNode(torch::lazy::OpKind{c10::Symbol::fromQualString("aten::ne")},
+              {a, b}, xla::ShapeUtil::MakeShape(xla::S64, {}), 1) {
+  const torch::lazy::DimensionNode* dim_node_0 = DimCast(operand(0));
+  const torch::lazy::DimensionNode* dim_node_1 = DimCast(operand(1));
+  XLA_CHECK(dim_node_0);
+  XLA_CHECK(dim_node_1);
+};
+
+int64_t SizeNe::getDynamicValue() const {
+  const torch::lazy::DimensionNode* dim_node_0 = DimCast(operand(0));
+  const torch::lazy::DimensionNode* dim_node_1 = DimCast(operand(1));
+  XLA_CHECK(dim_node_0);
+  XLA_CHECK(dim_node_1);
+  return dim_node_0->getDynamicValue() != dim_node_1->getDynamicValue() ? 1 : 0;
+}
+
+std::string SizeNe::ToString() const { return "SizeNe"; }
+
 SizeConstant::SizeConstant(int64_t val) : Scalar(c10::Scalar{val}, xla::S64){};
 
 SizeMul::SizeMul(torch::lazy::Value a, torch::lazy::Value b)

torch_xla/csrc/ops/dynamic_ir.h

@@ -68,6 +68,21 @@ class SizeEq : public XlaNode, public torch::lazy::DimensionNode {
   }
 };
 
+class SizeNe : public XlaNode, public torch::lazy::DimensionNode {
+ public:
+  SizeNe(torch::lazy::Value a, torch::lazy::Value b);
+  int64_t getDynamicValue() const override;
+  int64_t getStaticValue() const override {
+    TORCH_CHECK(false, "Comparison operators should be using getDynamicValue");
+  }
+  bool isSymbolic() const override { return true; }
+  std::string ToString() const override;
+  virtual XlaOpVector Lower(LoweringContext* loctx) const override {
+    // TODO: not sure we will ever need it?
+    TORCH_CHECK(false, "Lowering comparison nodes isn't supported yet!");
+  }
+};
+
 class SizeAdd : public XlaNode, public torch::lazy::DimensionNode {
  public:
   SizeAdd(torch::lazy::Value a, torch::lazy::Value b);

torch_xla/csrc/tensor.cpp

@@ -615,6 +615,12 @@ c10::SymNode XLASymNodeImpl::eq(const c10::SymNode& other) {
   return c10::make_intrusive<XLASymNodeImpl>(neq);
 }
 
+c10::SymNode XLASymNodeImpl::ne(const c10::SymNode& other) {
+  auto pother = dynamic_cast<XLASymNodeImpl*>(other.get());
+  auto nne = torch::lazy::MakeNode<SizeNe>(node(), pother->node());
+  return c10::make_intrusive<XLASymNodeImpl>(nne);
+}
+
 c10::SymNode XLASymNodeImpl::add(const c10::SymNode& other) {
   auto pother = dynamic_cast<XLASymNodeImpl*>(other.get());
   auto nadd = torch::lazy::MakeNode<SizeAdd>(node(), pother->node());

torch_xla/csrc/tensor.h

@@ -33,6 +33,7 @@ class TORCH_API XLASymNodeImpl : public c10::SymNodeImpl {
   bool is_int() override;
   bool is_float() override;
   c10::SymNode eq(const c10::SymNode& other) override;
+  c10::SymNode ne(const c10::SymNode& other) override;
   c10::SymNode add(const c10::SymNode& other) override;
   c10::SymNode mul(const c10::SymNode& other) override;
   c10::SymNode floordiv(const c10::SymNode& other) override;