Is there a unfold function in JAX?

[jingtianweidizzf]

Dear all,

thank you for your wonderful work! I have a question, is there a function as torch.nn.F.unfold which transforms the input images into matrix (or like the im2col function in pytorch c++ code.) already build in JAX?

What I can find now is a version built with numpy here, but it seems that the jnp.pad(X, ((0,0), (0,0), (pad, pad), (pad, pad))) function don't support the JIT compile? And if I build this function with numpy in python it is slower than the torch.nn.F.unfold.

Thank you for all information! :)

Best,
Zifan

[mattjj]

Thanks for the question, and the kind wodrs!

I think conv_general_dilated_patches might do what we want here. It's a pretty general function and there might be room for convenience wrappers.

What do you think?

[jingtianweidizzf]

Thanks for the reply:)

I'm not sure if that's what we want, I ran into a problem when calling that function.

Firstly, what I exactly want is like the following:

x1 = torch.arange(16.0).reshape(1,1,4,4)
print(x1)
x1_unf = F.unfold(x1,kernel_size = 3,padding = 0).transpose(1,2)
print(x1_unf)
tensor([[[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.],
          [12., 13., 14., 15.]]]])
tensor([[[ 0.,  1.,  2.,  4.,  5.,  6.,  8.,  9., 10.],
         [ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
         [ 4.,  5.,  6.,  8.,  9., 10., 12., 13., 14.],
         [ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.]]])

And now I'm trying to call that function by

x2 = jnp.arange(16.0).reshape([1,1,4,4])
x2_unf = jax.lax.conv_general_dilated_patches(lhs = x2, filter_shape = [3,3], 
                                                      window_strides = [1,1], padding = [0,0])

But the following error appears and I have no idea what to do next:

TypeError: ConvGeneralDilated(): incompatible function arguments. The following argument types are supported:
  1. (lhs: jaxlib.xla_extension.XlaOp, rhs: jaxlib.xla_extension.XlaOp, window_strides: Span[int], padding: Span[Tuple[int, int]], lhs_dilation: Span[int], rhs_dilation: Span[int], dimension_numbers: xla::ConvolutionDimensionNumbers, feature_group_count: int = 1, batch_group_count: int = 1, precision_config: xla::PrecisionConfig = None, preferred_element_type: Optional[jaxlib.xla_extension.PrimitiveType] = None) -> jaxlib.xla_extension.XlaOp

Invoked with: <jaxlib.xla_extension.XlaOp object at 0x7f0b6b3f0cb0>, <jaxlib.xla_extension.XlaOp object at 0x7f0b6b3f06b0>, (1, 1), (0, 0), (1, 1), (1, 1), <jaxlib.xla_client.ConvolutionDimensionNumbers object at 0x7f0b6b3c5250>, 1, 1; kwargs: precision_config=None

Do you have any examples of calling that function or any suggestion to solve this problem?

Best,
Zifan

[jingtianweidizzf]

Sorry for my stupid question again. I figured out why it didn't work and now it's fine.

[lyuxingjian]

@mattjj Hi! I was trying the solution and ran the following code:

import jax.numpy as jnp

dev = jax.local_devices()[0]

ftmap_fn = jax.device_put(jnp.ones((1, 384, 384, 256), dtype=jnp.bfloat16), device=dev)
W = 5
stride, padding = (4, 4), (W // 2, W // 2)
jax.lax.conv_general_dilated_patches(ftmap_fn, (5, 5), stride, padding).shape

but i found the following error:

...
File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:609, in XlaCompiledComputation.from_xla_computation(name, xla_computation, nreps, device, backend, tuple_args, in_avals, out_avals, kept_var_idx)
    606 options.parameter_is_tupled_arguments = tuple_args
    607 with log_elapsed_time(f"Finished XLA compilation of {name} "
    608                       "in {elapsed_time} sec"):
--> 609   compiled = compile_or_get_cached(backend, xla_computation, options)
    610 buffer_counts = (None if len(out_avals) == 1 else
    611                  [aval_to_num_buffers(aval) for aval in out_avals])
    612 execute = _execute_compiled if nreps == 1 else _execute_replicated

File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:578, in compile_or_get_cached(backend, computation, compile_options)
    575   ir_str = (computation if isinstance(computation, str)
    576             else computation.as_hlo_text())
    577   _dump_ir_to_file(module_name, ir_str)
--> 578 return backend_compile(backend, computation, compile_options)

File ~/.local/lib/python3.8/site-packages/jax/_src/profiler.py:206, in annotate_function.<locals>.wrapper(*args, **kwargs)
    203 @wraps(func)
    204 def wrapper(*args, **kwargs):
    205   with TraceAnnotation(name, **decorator_kwargs):
--> 206     return func(*args, **kwargs)
    207   return wrapper

File ~/.local/lib/python3.8/site-packages/jax/_src/dispatch.py:532, in backend_compile(backend, built_c, options)
    528 @profiler.annotate_function
    529 def backend_compile(backend, built_c, options):
    530   # we use a separate function call to ensure that XLA compilation appears
    531   # separately in Python profiling results
--> 532   return backend.compile(built_c, compile_options=options)

RuntimeError: UNKNOWN: -:4:130: error: expected '['

The error looks very cryptic cannot figure out what it means. Is it an XLA / Jax bug? Thanks!

[jakevdp]

This looks like it may be a bug – please open a new issue to report it, and please include information about the version of jax and jaxlib you're using, the OS/backend, and any other relevant information to help locate the source of the bug. Thanks!

[oligologarithmos]

sorry for necroposting but is there a backward operation for conv_general_dilated_patches, equivalent to fold

[mattjj]

Can you just apply jax.vjp or jax.linear_transpose to it? (Sorry, I’m afk, otherwise I would write a code example.)