Enhance torch.zeros_like for GPU inference in dynamic shape

coremltools/converters/mil/frontend/torch/ops.py

@@ -5529,23 +5529,29 @@ def zeros_like(context, node):
  min_expected={TorchFrontend.TORCHEXPORT: 1, TorchFrontend.EXECUTORCH: 1},
  )
  x = inputs[0]
- shape = mb.shape(x=x)
+ src_np_type = nptype_from_builtin(x.dtype)
  if len(inputs) > 1 and inputs[1] and inputs[1].val:
  dtype = inputs[1].val
- np_type = NUM_TO_NUMPY_DTYPE[dtype]
+ dst_np_type = NUM_TO_NUMPY_DTYPE[dtype]
  else:
- np_type = nptype_from_builtin(x.dtype)
+ dst_np_type = src_np_type
 
+ shape = mb.shape(x=x)
  if shape.can_be_folded_to_const():
  shape = shape.val
- zeros = _np.zeros(shape).astype(np_type)
+ zeros = _np.zeros(shape).astype(dst_np_type)
  zeros_like = mb.const(val=zeros, name=node.name)
  else:
- value = np_type(0)
- if is_current_opset_version_compatible_with(target.iOS16):
- zeros_like = mb.fill_like(ref_tensor=x, value=value, name=node.name)
+ if src_np_type == np.bool_:
+ zeros = mb.logical_xor(x=x, y=x)
+ else:
+ zeros = mb.sub(x=x, y=x)
+ if src_np_type != dst_np_type:
+ num = NUMPY_DTYPE_TO_TORCH_NUM[dst_np_type]
+ zeros_like = mb.cast(x=zeros, dtype=NUM_TO_DTYPE_STRING[num], name=node.name)
  else:
- zeros_like = mb.fill(shape=shape, value=value, name=node.name)
+ zeros.name = node.name
+ zeros_like = zeros
 
  context.add(zeros_like)
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -7402,6 +7402,34 @@ def forward(self, x):
  # The zeros op is not folded to const.
  assert len(prog.find_ops(op_type="fill")) == 1
 
+ @pytest.mark.parametrize(
+ "compute_unit, backend, is_dynamic, src_dtype, dst_dtype",
+ itertools.product(
+ compute_units,
+ backends,
+ [True, False],
+ [torch.float16, torch.float32, torch.int32, torch.bool],
+ [torch.float16, torch.float32, torch.int32, torch.bool],
+ ),
+ )
+ def test_zeros_like_types(self, compute_unit, backend, is_dynamic, src_dtype, dst_dtype):
+ target, type = None, None
+ if src_dtype == torch.float16 or dst_dtype == torch.float16:
+ target = ct.target.iOS16
+ if is_dynamic:
+ type = [ct.TensorType(shape=ct.Shape([ct.RangeDim(1, 1_000)]))]
+ model = ModuleWrapper(function=torch.zeros_like, kwargs={"dtype": dst_dtype}).eval()
+
+ self.run_compare_torch(
+ torch.tensor([3], dtype=src_dtype),
+ model,
+ input_as_shape=False,
+ backend=backend,
+ compute_unit=compute_unit,
+ converter_input_type=type,
+ minimum_deployment_target=target,
+ )
+
 
 class TestTopk(TorchBaseTest):
  @pytest.mark.parametrize(