Example of applying xmap on serial

[mattiasmar]

Looking for an example applying xmap on the elements of a stax.serial
Thanks!

[hawkinsp]

This doesn't answer your question, but please note: we discourage the use of stax for most users who are looking for a neural network library. In general, we point users to Flax or Haiku.

Can you say more about what you are trying to do? I can't entirely parse the question.

[mattiasmar]

I created the MLP example below in an attempt to control the pjit of individual layers. Could you please help me define the tuple of PartitionSpec's of "Example 2"?

import jax
import jax.numpy as jnp
from jax import random
from jax.experimental.pjit import pjit
from jax.experimental import PartitionSpec, maps
from jax.scipy.special import logsumexp
import numpy as np
# A helper function to randomly initialize weights and biases
# for a dense neural network layer
def random_layer_params(n_in, n_out, key, scale=1e-2):
    w_key, b_key = random.split(key)
    return scale * random.normal(w_key, (n_in, n_out)), scale * random.normal(b_key, (n_out,))

# Initialize all layers for a fully-connected neural network with sizes "sizes"
def init_network_params(sizes, key):
    keys = random.split(key, len(sizes))
    return [random_layer_params(n_in, n_out, k) for n_in, n_out, k in zip(sizes[:-1], sizes[1:], keys)]

def relu(x):
    return jnp.maximum(0,x)

params = init_network_params([8,16,32,10], random.PRNGKey(0))
def predict(params, image):
  # per-example predictions
    activations = image
    for w, b in params[:-1]:
        outputs = jnp.dot(activations,w) + b
        activations = relu(outputs)
  
    final_w, final_b = params[-1]
    logits = jnp.dot(activations,final_w ) + final_b
    return logits - logsumexp(logits)
input = np.repeat(np.random.randint(0,9,(1,8)),repeats=4,axis=0)
mesh = maps.Mesh(np.array(jax.devices()).reshape(2,2), ('x','y'))

v,t = jax.tree_flatten(params)
print(t) 
#=========
# PyTreeDef([(*, *), (*, *), (*, *)])
#=========

#==========Example 1: WORKS=============
print("Basic pjit")
f = pjit( predict,in_axis_resources=PartitionSpec('x'),out_axis_resources=None)
with maps.mesh(mesh.devices, mesh.axis_names):
    out = f(params,input)



#==========Example 2:  DOESN'T WORK=============
print("Per layer pjit")
f = pjit( predict,in_axis_resources=(PartitionSpec('x'),PartitionSpec('x'),PartitionSpec('x')), out_axis_resources=None)
with maps.mesh(mesh.devices, mesh.axis_names):
    out = f(params,input)

#ValueError: pjit in_axis_resources specification must be a tree prefix of the corresponding value, 
#got specification (PartitionSpec('x',), PartitionSpec('x',), PartitionSpec('x',)) for value tree 
#PyTreeDef(([(*, *), (*, *), (*, *)], *)). Note that pjit in_axis_resources that are non-trivial 
#pytrees should always be wrapped in a tuple representing the argument list.

[skye]

GpuDevices never have a coords field, because they're generally not arranged in a mesh topology, so this behavior is expected. I believe the coords are used in Marc's example to generate a topology-aware device mesh, so as a workaround, you could add something like this to the top of get_mesh in siren.py:

if jax.default_backend() == 'gpu':
  return maps.Mesh(np.array(jax.devices()).reshape(2,2), ('data','model'))
assert jax.default_backend() == 'tpu'

[skye]

For your example 2 that doesn't work, you're providing 3 PartitionSpecs, but your function f only takes two arguments.

[mattiasmar]

@skye Thanks. The intention behind example 2 was that the 3 PartitionSpecs would match the 3 layers listed by tree_flatten: (PyTreeDef([(*, *), (*, *), (*, *)])) and eventually I want to give a different PartionSpec for each layer. From what I can tell from the documentation the argument in_axis_resources and out_axis_resources of pjit must match (or be broadcasted) to the PyTreeDef of the input fun to pjit (in our case the argument params of method predict).

Could you tell me what I got wrong in my understanding (or could you help me achieve this layer wise control of PartionSpec that I'm trying to achieve)?

[mattiasmar]

If I decorate each of the layers inside the predict function with pjit (and accompany each such decoration with a PartionSpec at the same time) it works (obviously).
However, my goal is to decorate all of the layers of the model (ideally any model) in a post-processing step. This is in contrast to surgically entering the code and adding calls to pjit in attachment to every layer (which can be prohibitively complicated to do).

So if I somehow can give a list of PartionSpec's (or a tuple/dict/PyTree of PartionSpec's) that relates to all the layers of model in a single pjit call, then I would have achieved my goal. Is this possible? Alternatively is this possible with xmap?

[mattiasmar]

I an attempt to be even more specific, and to communicate to you better, I defined the params object as a dictionary with keys ("weights" and "biases")
Inside the code I added the print outs, so you don't actually need to run the code.

import jax
import jax.numpy as jnp
from jax import random, pmap
from jax.experimental.pjit import pjit
from jax.experimental import PartitionSpec, maps
from jax.scipy.special import logsumexp
import numpy as np
import os

os.environ['NCCL_LAUNCH_MODE'] = "PARALLEL"

# A helper function to randomly initialize weights and biases
# for a dense neural network layer
def random_layer_params(n_in, n_out, key, scale=1e-2):
    w_key, b_key = random.split(key)
    return {"weights": scale * random.normal(w_key, (n_in, n_out)), "biases": scale * random.normal(b_key, (n_out,))}


# Initialize all layers for a fully-connected neural network with sizes "sizes"
def init_network_params(sizes, key):
    keys = random.split(key, len(sizes))
    return [random_layer_params(n_in, n_out, k) for n_in, n_out, k in zip(sizes[:-1], sizes[1:], keys)]

def relu(x):
    return jnp.maximum(0, x)

layers = [28 * 28,512,10]
params = init_network_params(layers, random.PRNGKey(0))

def f(activations, layer):
    outputs = jnp.dot(activations, layer['weights']) + layer['biases']
    activations = relu(outputs)
    return activations


def predict(params, image):
    # per-example predictions
    activations = image
    for i, layer in enumerate(params[:-1]):
        activations = f(activations, layer)

    logits = f(activations, params[-1])
    return logits - logsumexp(logits)


input = np.repeat(np.random.randint(0, 9, (1, 28 * 28)), repeats=8 * 256, axis=0)
mesh = maps.Mesh(np.array(jax.devices()).reshape(2, 4), ('x', 'y'))

v, t = jax.tree_flatten(params)
print("Structure of 'params'=",t)
#=========
#Structure of 'params'= PyTreeDef([{'biases': *, 'weights': *}, {'biases': *, 'weights': *}])
#=========

specs= list()
for _ in range(len(layers)-1):
  specs.append({"biases":PartitionSpec('x'),"weights":PartitionSpec('x')})
specs.append(PartitionSpec(None)) #NOTE! Both with and without this line the error message remains the same.
specs = tuple(specs)
print("Structure of 'PartitionSpec Tree'=", specs)
#==========
# Structure of 'PartitionSpec Tree'= ({'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, {'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, PartitionSpec(None,))
#==========

f1 = pjit( predict,in_axis_resources=specs, out_axis_resources=None)
#=======
# FAILURE:
# ValueError: pjit in_axis_resources specification must be a tree prefix of the corresponding value,
#  got specification         ({'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, {'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, PartitionSpec(None,))
#  for value tree PyTreeDef(([{'biases': *,                   'weights':                   *}, {'biases': *,                   'weights':                  *}],                   *)).
#  Note that pjit in_axis_resources that are non-trivial pytrees should always be wrapped in a tuple representing the argument list.
#=======
with maps.mesh(mesh.devices, mesh.axis_names):
     out = f1(params,input)
print(out.shape)

[mattiasmar]

OK! Now I got it right! 😄
This is the story: Under the hood the jax wraps method (of utils.py) adds an additional element to the PyTreeDef. When accounting for that my code runs without errors (yet to determine if the math is preserved too).
Thanks @skye @avital @hawkinsp @marcvanzee

[skye]

Hi @mattiasmar , I'm not sure what you mean by wraps adding an additional elements. In your new example, you should indeed add the extra PartitionSpec(None) for the image argument of predict. However, you need to put all the layer PartitionSpecs in a list, so it matches the list structure of the params argument. Here's a cleaned up version of the error message where you can see that the specification is the same as the value PyTreeDef except for the missing list:

got specification (
{'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, 
{'biases': PartitionSpec('x',), 'weights': PartitionSpec('x',)}, 
PartitionSpec(None,)
)

for value tree PyTreeDef( (
[{'biases': *, 'weights': *}, 
 {'biases': *, 'weights': *}], 
*)
).

I think the following code should work:

layer_specs= list()
for _ in range(len(layers)-1):
  layer_specs.append({"biases":PartitionSpec('x'),"weights":PartitionSpec('x')})
image_spec = PartitionSpec(None)
specs = (layer_specs, image_spec)
print("Structure of 'PartitionSpec Tree'=", specs)

[mattiasmar]

Clearly the second PartitionSpec is needed for "image" argument of predict. Thanks for pointing that out @skye.