Add support for trig, binary, algebraic ops

experimental/torch_xla2/test/test_ops.py

@@ -155,7 +155,6 @@
  "logcumsumexp",
  "logdet",
  "logspace",
- "logsumexp",
  "lu",
  "lu_solve",
  "lu_unpack",
@@ -322,7 +321,6 @@
  "squeeze",
  "stft",
  "sub",
- "sum",
  "svd",
  "svd_lowrank",
  "take_along_dim",
@@ -352,14 +350,6 @@
  "trapz",
  "H",
  "T",
- "__radd__",
- "__rand__",
- "__rdiv__",
- "__rmod__",
- "__rmul__",
- "__ror__",
- "__rsub__",
- "__rxor__",
  "_softmax_backward_data",
  "addcdiv",
  "addcmul",
@@ -369,12 +359,7 @@
  "atleast_2d",
  "atleast_3d",
  "bfloat16",
- "bitwise_and",
- "bitwise_not",
- "bitwise_or",
- "bitwise_xor",
  "bmm",
- "bool",
  "broadcast_shapes",
  "cartesian_prod",
  "cdouble",
@@ -391,16 +376,12 @@
  "conj",
  "conj_physical",
  "contiguous",
- "cos",
- "cosh",
  "deg2rad",
  "diagonal",
  "dot",
- "double",
  "dsplit",
  "dstack",
  "einsum",
- "eq",
  "erf",
  "exp",
  "expm1",
@@ -412,57 +393,28 @@
  "flip",
  "fliplr",
  "flipud",
- "float",
- "float_power",
- "floor",
  "fmod",
  "frac",
  "full",
- "ge",
  "gradient",
- "gt",
- "half",
  "hsplit",
  "hstack",
  "index_add",
  "inner",
- "int",
- "isfinite",
- "isinf",
- "isnan",
- "isneginf",
- "isposinf",
- "isreal",
  "kron",
  "ldexp",
- "le",
  "linalg.diagonal",
  "linalg.vecdot",
- "log10",
- "log1p",
- "log2",
- "log",
- "logical_and",
- "logical_not",
- "logical_or",
- "logical_xor",
- "logit",
- "long",
- "lt",
  "mH",
  "mT",
  "masked_fill",
- "maximum",
  "meshgrid",
- "minimum",
  "movedim",
  "msort",
  "mul",
  "mv",
  "nan_to_num",
  "native_dropout_backward",
- "ne",
- "neg",
  "new_full",
  "new_ones",
  "new_zeros",
@@ -501,11 +453,9 @@
  "outer",
  "permute",
  "positive",
- "pow",
  "rad2deg",
  "randint",
  "ravel",
- "real",
  "reciprocal",
  "remainder",
  "repeat",
@@ -514,13 +464,8 @@
  "resolve_conj",
  "resolve_neg",
  "roll",
- "round",
  "rsqrt",
  "scalar_tensor",
- "sgn",
- "short",
- "sigmoid",
- "sign",
  "signal.windows.bartlett",
  "signal.windows.blackman",
  "signal.windows.cosine",
@@ -532,9 +477,6 @@
  "signal.windows.hann",
  "signal.windows.kaiser",
  "signal.windows.nuttall",
- "sin",
- "sinc",
- "sinh",
  "slice",
  "slice_scatter",
  "sparse.mm",
@@ -545,19 +487,11 @@
  "split",
  "split_with_sizes",
  "split_with_sizes_copy",
- "sqrt",
  "square",
  "stack",
- "std",
- "std_mean",
  "sum_to_size",
- "t",
- "tan",
- "tanh",
  "tensordot",
  "tile",
- "to",
- "transpose",
  "tril",
  "tril_indices",
  "triu_indices",
@@ -567,16 +501,13 @@
  "unsafe_chunk",
  "unsafe_split",
  "unsqueeze",
- "var",
  "vdot",
  "view_as_complex",
  "view_as",
  "view_copy",
  "view",
  "vsplit",
  "vstack",
- "where",
- "xlogy",
  "zeros",
 }
 

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -911,6 +911,8 @@ def _aten_bitwise_not(self):
 def _aten_sum(self, dim=None, keepdim=False, dtype=None):
  if not dim:
  dim = None
+ if self.ndim == 0:
+ return self
  return jnp.sum(self, axis=dim, keepdims=keepdim, dtype=dtype)
 
 
@@ -1447,8 +1449,12 @@ def _aten_diagonal(input, offset=0, dim1=0, dim2=1):
 # aten.eq
 @op(torch.ops.aten.eq)
 def _aten_eq(input1, input2):
- return input1 == input2
+ return jnp.equal(input1, input2)
 
+# aten.equal
+@op(torch.ops.aten.equal)
+def _aten_equal(input1, input2):
+ return input1 == input2
 
 # aten.erf
 @op(torch.ops.aten.erf)