support log2 (apple#1568)

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

@@ -4519,6 +4519,13 @@ def log10(context, node):
     log_x = mb.log(x=x)
     context.add(mb.mul(x=log_x, y=1/_np.log(10.0)), node.name)
 
+@register_torch_op
+def log2(context, node):
+    inputs = _get_inputs(context, node)
+    x = inputs[0]
+    log_x = mb.log(x=x)
+    context.add(mb.mul(x=log_x, y=1/_np.log(2.0)), node.name)
+
 @register_torch_op
 def flip(context, node):
     inputs = _get_inputs(context, node, expected=2)

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

@@ -3462,6 +3462,25 @@ def forward(self, x):
             input_shape, model, backend=backend,
         )
 
+class TestLog2(TorchBaseTest):
+    @pytest.mark.parametrize(
+        "backend, rank", itertools.product(backends, range(1, 6)),
+    )
+    def test_log2(self, backend, rank):
+
+        class Log2Model(nn.Module):
+            def __init__(self):
+                super(Log2Model, self).__init__()
+
+            def forward(self, x):
+                return torch.log2(x)
+
+        input_shape = tuple(np.random.randint(low=1, high=10, size=rank))
+        model = Log2Model()
+        self.run_compare_torch(
+            input_shape, model, backend=backend,
+        )
+
 class TestFlip(TorchBaseTest):
     @pytest.mark.parametrize(
         "backend, rank_dim",