feat(atenlib): where, full (#286)

Implement where, full, new_full, full_like

Author: justinchuby

onnxscript/function_libs/torch_aten/ops/core.py

@@ -1841,18 +1841,23 @@ def aten_from_file(
     raise NotImplementedError()
 
 
-def aten_full(size: INT64, fill_value: float) -> TensorType:
+@torch_op("aten::full")
+def aten_full(size: INT64, fill_value, dtype: int = FLOAT.dtype):
     # full(SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
 
-    raise NotImplementedError()
+    fill_value = op.Cast(fill_value, to=dtype)
 
+    return op.Expand(fill_value, size)
 
-def aten_full_like(
-    self: TensorType, fill_value: float, memory_format: Optional[str] = None
-) -> TensorType:
+
+@torch_op("aten::full_like")
+def aten_full_like(self, fill_value, dtype: int = FLOAT.dtype):
     # full_like(Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor
 
-    raise NotImplementedError()
+    fill_value = op.Cast(fill_value, to=dtype)
+    self_shape = op.Shape(self)
+
+    return op.Expand(fill_value, self_shape)
 
 
 def aten_fused_moving_avg_obs_fake_quant(
@@ -3447,10 +3452,15 @@ def aten_new_empty_strided(self: TensorType, size: INT64, stride: INT64) -> Tens
     raise NotImplementedError()
 
 
-def aten_new_full(self: TensorType, size: INT64, fill_value: float) -> TensorType:
+@torch_op("aten::new_full")
+def aten_new_full(
+    self, size: INT64, fill_value, dtype: int = FLOAT.dtype
+):  # pylint: disable=unused-argument
     # new_full(Tensor self, SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor
 
-    raise NotImplementedError()
+    fill_value = op.Cast(fill_value, to=dtype)
+
+    return op.Expand(fill_value, size)
 
 
 def aten_new_ones(self: TensorType, size: INT64) -> TensorType:
@@ -4928,10 +4938,11 @@ def aten_vstack(tensors: Sequence[TensorType]) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_where(condition: TensorType) -> TensorType:
-    # where(Tensor condition) -> Tensor[]
+@torch_op("aten::where")
+def aten_where(self: TTensor, condition: BOOL, other: TTensor) -> TTensor:
+    # where.self(Tensor condition, Tensor self, Tensor other) -> Tensor
 
-    raise NotImplementedError()
+    return op.Where(condition, self, other)
 
 
 def aten_xlogy(self: TensorType, other: TensorType) -> TensorType:

onnxscript/test/function_libs/torch_aten/ops_correctness_test.py

@@ -169,6 +169,8 @@ def wrapped(fn):
     "exp": core_ops.aten_exp,
     "exp2": core_ops.aten_exp2,
     "fmod": core_ops.aten_fmod,
+    # TODO(justinchuby): Test aten::full
+    "full_like": core_ops.aten_full_like,
     "gt": core_ops.aten_gt,
     "isinf": core_ops.aten_isinf,
     "lt": core_ops.aten_lt,
@@ -177,6 +179,7 @@ def wrapped(fn):
     "mul": core_ops.aten_mul,
     "ne": core_ops.aten_ne,
     "neg": core_ops.aten_neg,
+    "new_full": core_ops.aten_new_full,
     "nn.functional.elu": nn_ops.aten_elu,
     "nn.functional.leaky_relu": nn_ops.aten_leaky_relu,
     "nn.functional.linear": nn_ops.aten_linear,
@@ -202,6 +205,7 @@ def wrapped(fn):
     "tanh": core_ops.aten_tanh,
     # "transpose": core_ops.aten_transpose,  # TODO(justinchuby): Enable when onnxscript errors are fixed,
     "unsqueeze": core_ops.aten_unsqueeze,
+    "where": core_ops.aten_where,
     "zeros": core_ops.aten_zeros,
     "zeros_like": core_ops.aten_zeros_like,
 }