<p dir="auto">Should it be possible to define a network in stax without setting a non-batch dimension?</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as np
from jax import random
from jax.experimental import stax
from jax.experimental.stax import Dense, Relu, Flatten, LogSoftmax

# Use stax to set up network initialization and evaluation functions
net_init, net_apply = stax.serial(Dense(128), Relu, Dense(10), LogSoftmax)

rng = random.PRNGKey(0)
# Initialize parameters, not committing to a batch shape

# works
in_shape = (-1, 1)
out_shape, net_params = net_init(rng, in_shape)

# doesn't work (fails with ValueError: all elements of broadcast shape must be non-negative)
in_shape = (-1,)
out_shape, net_params = net_init(rng, in_shape)"><pre class="notranslate"><code class="notranslate">import jax.numpy as np
from jax import random
from jax.experimental import stax
from jax.experimental.stax import Dense, Relu, Flatten, LogSoftmax

# Use stax to set up network initialization and evaluation functions
net_init, net_apply = stax.serial(Dense(128), Relu, Dense(10), LogSoftmax)

rng = random.PRNGKey(0)
# Initialize parameters, not committing to a batch shape

# works
in_shape = (-1, 1)
out_shape, net_params = net_init(rng, in_shape)

# doesn't work (fails with ValueError: all elements of broadcast shape must be non-negative)
in_shape = (-1,)
out_shape, net_params = net_init(rng, in_shape)
</code></pre></div>
<p dir="auto">As a workaround, I can of course add a fake 1 dimension on the end, and then remove it later, but I was a little surprised that this doesn't automatically work.</p>
<p dir="auto">The following code snippet shows that Stax MLPs can be defined w.r.t. unbatched examples (input_size = (1,)) while Convnets seem to require a batch size (though it can be -1). Is this intended behavior?</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="# Works
net_init, net_apply = stax.serial(
    Conv(32, (3, 3), padding='SAME'), Relu,
    Conv(64, (3, 3), padding='SAME'), Relu,
    MaxPool((2, 2)), Flatten,
    Dense(128), Relu,
    Dense(10), LogSoftmax,
)

# Initialize parameters, not committing to a batch shape
in_shape = (-1, 28, 28, 1)
out_shape, net_params = net_init(in_shape)

# Works
net_init, net_apply = stax.serial(
    Dense(40), Relu,
    Dense(40), Relu,
    Dense(1)
)
in_shape = (1,)
out_shape, net_params = net_init(in_shape)

# Doesn't Work
net_init, net_apply = stax.serial(
    Conv(32, (3, 3), padding='SAME'), Relu,
    Conv(64, (3, 3), padding='SAME'), Relu,
    MaxPool((2, 2)), Flatten,
    Dense(128), Relu,
    Dense(10), LogSoftmax,
)
in_shape = (28, 28, 1)
out_shape, net_params = net_init(in_shape)"><pre class="notranslate"><code class="notranslate"># Works
net_init, net_apply = stax.serial(
    Conv(32, (3, 3), padding='SAME'), Relu,
    Conv(64, (3, 3), padding='SAME'), Relu,
    MaxPool((2, 2)), Flatten,
    Dense(128), Relu,
    Dense(10), LogSoftmax,
)

# Initialize parameters, not committing to a batch shape
in_shape = (-1, 28, 28, 1)
out_shape, net_params = net_init(in_shape)

# Works
net_init, net_apply = stax.serial(
    Dense(40), Relu,
    Dense(40), Relu,
    Dense(1)
)
in_shape = (1,)
out_shape, net_params = net_init(in_shape)

# Doesn't Work
net_init, net_apply = stax.serial(
    Conv(32, (3, 3), padding='SAME'), Relu,
    Conv(64, (3, 3), padding='SAME'), Relu,
    MaxPool((2, 2)), Flatten,
    Dense(128), Relu,
    Dense(10), LogSoftmax,
)
in_shape = (28, 28, 1)
out_shape, net_params = net_init(in_shape)
</code></pre></div>
<p dir="auto">The last one returns the following error:<br>
IndexError                                Traceback (most recent call last)<br>
 in ()<br>
9 # Initialize parameters, not committing to a batch shape<br>
10 in_shape = (28, 28, 1)<br>
---&gt; 11 out_shape, net_params = net_init(in_shape)</p>
<p dir="auto">google3/third_party/py/jax/experimental/stax.py in init_fun(input_shape)<br>
269     params = []<br>
270     for init_fun in init_funs:<br>
--&gt; 271       input_shape, param = init_fun(input_shape)<br>
272       params.append(param)<br>
273     return input_shape, params</p>
<p dir="auto">google3/third_party/py/jax/experimental/stax.py in init_fun(input_shape)<br>
109     kernel_shape = [out_chan if c == 'O' else<br>
110                     input_shape[lhs_spec.index('C')] if c == 'I' else<br>
--&gt; 111                     next(filter_shape_iter) for c in rhs_spec]<br>
112     output_shape = lax.conv_general_shape_tuple(<br>
113         input_shape, kernel_shape, strides, padding, dimension_numbers)</p>
<p dir="auto">IndexError: tuple index out of range</p>