Add tf.nn.conv to the JAX substrate.

Also, remove another line adjusting rewrite from the NumPy rewrite system.

PiperOrigin-RevId: 396021511

Author: SiegeLordEx

tensorflow_probability/python/internal/backend/jax/rewrite.py

@@ -45,7 +45,7 @@ def main(argv):
   contents = contents.replace('scipy.special', 'jax.scipy.special')
   if FLAGS.rewrite_numpy_import:
     contents = contents.replace('\nimport numpy as np',
-                                '\nimport numpy as onp\nimport jax.numpy as np')
+                                '\nimport numpy as onp; import jax.numpy as np')
   else:
     contents = contents.replace('\nimport numpy as np',
                                 '\nimport numpy as np; onp = np')

tensorflow_probability/python/internal/backend/numpy/nn.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
+
 # Dependency imports
 import numpy as np
 
@@ -35,6 +37,7 @@
 
 
 __all__ = [
+    'conv2d',
     'l2_normalize',
     'log_softmax',
     'moments',
@@ -48,6 +51,106 @@
 ]
 
 
+JAX_MODE = False
+
+
+if JAX_MODE:
+  import jax  # pylint: disable=g-import-not-at-top
+
+
+# Borrowed from TensorFlow.
+def _get_sequence(value, n, channel_index, name):
+  """Formats a value input for gen_nn_ops."""
+  # Performance is fast-pathed for common cases:
+  # `None`, `list`, `tuple` and `int`.
+  if value is None:
+    return [1] * (n + 2)
+
+  # Always convert `value` to a `list`.
+  if isinstance(value, list):
+    pass
+  elif isinstance(value, tuple):
+    value = list(value)
+  elif isinstance(value, int):
+    value = [value]
+  elif not isinstance(value, collections.abc.Sized):
+    value = [value]
+  else:
+    value = list(value)  # Try casting to a list.
+
+  len_value = len(value)
+
+  # Fully specified, including batch and channel dims.
+  if len_value == n + 2:
+    return value
+
+  # Apply value to spatial dims only.
+  if len_value == 1:
+    value = value * n  # Broadcast to spatial dimensions.
+  elif len_value != n:
+    raise ValueError('{} should be of length 1, {} or {} but was {}'.format(
+        name, n, n + 2, len_value))
+
+  # Add batch and channel dims (always 1).
+  if channel_index == 1:
+    return [1, 1] + value
+  else:
+    return [1] + value + [1]
+
+
+def _conv2d(
+    input,  # pylint: disable=redefined-builtin
+    filters,
+    strides,
+    padding,
+    data_format='NHWC',
+    dilations=None,
+    name=None):
+  """tf.nn.conv2d implementation."""
+  del name
+  if not JAX_MODE:
+    raise NotImplementedError('tf.nn.conv2d not implemented in NumPy.')
+
+  if dilations is not None:
+    raise ValueError('Dilations not yet supported')
+
+  channel_index = 1 if data_format.startswith('NC') else 3
+
+  window_strides = _get_sequence(strides, 2, channel_index, 'strides')
+  if window_strides[0] != 1:
+    raise ValueError(
+        f'Stride != 1 not supported for batch dimension. `strides`: {strides} '
+        f'`data_format`: {data_format}')
+  if window_strides[channel_index] != 1:
+    raise ValueError(
+        'Stride != 1 not supported for channel dimension. `strides`: '
+        f'{strides} '
+        f'`data_format`: {data_format}')
+  window_strides = [
+      e for i, e in enumerate(window_strides) if i not in [0, channel_index]
+  ]
+
+  if isinstance(padding, (list, tuple)):
+    if padding[0] != (0, 0):
+      raise ValueError(
+          f'Padding not supported for batch dimension. `padding`: {padding} '
+          f'`data_format`: {data_format}')
+    if padding[channel_index] != (0, 0):
+      raise ValueError(
+          'Padding not supported for channel dimension. `padding`: '
+          f'{padding} '
+          f'`data_format`: {data_format}')
+    padding = [s for i, s in enumerate(padding) if i not in [0, channel_index]]
+
+  return jax.lax.conv_general_dilated(
+      lhs=input,
+      rhs=filters,
+      window_strides=window_strides,
+      padding=padding,
+      dimension_numbers=(data_format, 'HWIO', data_format),
+  )
+
+
 def _sigmoid_cross_entropy_with_logits(  # pylint: disable=invalid-name,unused-argument
     _sentinel=None,
     labels=None,
@@ -96,6 +199,11 @@ def _softmax_cross_entropy_with_logits(  # pylint: disable=invalid-name,unused-a
 
 # --- Begin Public Functions --------------------------------------------------
 
+conv2d = utils.copy_docstring(
+    'tf.nn.conv2d',
+    _conv2d)
+
+
 l2_normalize = utils.copy_docstring(
     'tf.nn.l2_normalize',
     l2_normalize)

tensorflow_probability/python/internal/backend/numpy/numpy_test.py

@@ -616,6 +616,54 @@ def argsort_params(draw):
        True))  # stable sort
 
 
+@hps.composite
+def conv2d_params(draw):
+  # NCHW is GPU-only
+  # data_format = draw(hps.sampled_from(['NHWC', 'NCHW']))
+  data_format = draw(hps.just('NHWC'))
+
+  input_shape = draw(shapes(4, 4, min_side=5, max_side=10))
+  if data_format.startswith('NC'):
+    channels = input_shape[1]
+  else:
+    channels = input_shape[3]
+  filter_shape = draw(shapes(3, 3, min_side=1, max_side=4))
+  filter_shape = filter_shape[:2] + (channels, filter_shape[-1])
+
+  input_ = draw(
+      single_arrays(
+          batch_shape=(),
+          shape=hps.just(input_shape),
+      ))
+  filters = draw(single_arrays(
+      batch_shape=(),
+      shape=hps.just(filter_shape),
+  ))
+  small = hps.integers(0, 5)
+  small_pos = hps.integers(1, 5)
+  strides = draw(hps.one_of(small_pos, hps.tuples(small_pos, small_pos)))
+  if isinstance(strides, tuple) and len(strides) == 2 and draw(hps.booleans()):
+    if data_format.startswith('NC'):
+      strides = (1, 1) + strides
+    else:
+      strides = (1,) + strides + (1,)
+
+  zeros = (0, 0)
+  explicit_padding = (
+      draw(hps.tuples(small, small)),
+      draw(hps.tuples(small, small)),
+  )
+  if data_format.startswith('NC'):
+    explicit_padding = (zeros, zeros) + explicit_padding
+  else:
+    explicit_padding = (zeros,) + explicit_padding + (zeros,)
+  padding = draw(
+      hps.one_of(
+          hps.just(explicit_padding), hps.sampled_from(['SAME', 'VALID'])))
+
+  return (input_, filters, strides, padding, data_format)
+
+
 @hps.composite
 def sparse_xent_params(draw):
   num_classes = draw(hps.integers(1, 6))
@@ -1223,6 +1271,7 @@ def _not_implemented(*args, **kwargs):
              [n_same_shape(n=2, elements=[floats(), positive_floats()])]),
     TestCase('math.xlog1py',
              [n_same_shape(n=2, elements=[floats(), positive_floats()])]),
+    TestCase('nn.conv2d', [conv2d_params()], disabled=NUMPY_MODE),
     TestCase(
         'nn.sparse_softmax_cross_entropy_with_logits', [sparse_xent_params()],
         rtol=1e-4,
@@ -1276,9 +1325,8 @@ def _not_implemented(*args, **kwargs):
         xla_disabled=True),
     TestCase('histogram_fixed_width_bins',
              [histogram_fixed_width_bins_params()]),
-    TestCase(
-        'argsort', [argsort_params()],
-        xla_const_args=(1, 2, 3)),  # axis, direction, stable-sort
+    TestCase('argsort', [argsort_params()],
+             xla_const_args=(1, 2, 3)),  # axis, direction, stable-sort
 ]