Move signature implementation to a single function, and catch shape errors #35

phinate · 2023-07-12T12:34:50Z

I'm playing around with dispatching signax.signature in such a way that the correct behaviour happens depending on the shape of the path.

As an example, the function will still work on paths of shape (a, b ,c) and just (d) too, but both of these may not be the desired behaviour. So, this makes sure that if you feed in a path that is not right, we catch it as an error.

Also, if the path has a batch dim, this will broadcast the signature over the batch with vmap! flatten and num_chunks have become arguments to signature.

anh-tong · 2023-07-13T05:41:20Z

src/signax/signatures.py

-    return carry
+    else:
+        res = carry
+    if flatten:


Since the shape of res when stream=True (shape = [(path_len-1, dim), (path_len-1, dim, dim]) is different to when stream=False (shape = [(dim,), (dim, dim)]), we may use jax.vmap(utils.flatten) when stream=True.

We vmap the first dimesion path_len - 1

covered this by using ravel_pytree and getting rid of flatten:

import jax.numpy as jnp from jax import flatten_util arr = [jnp.ones((2, 3)), jnp.ones((2, 3, 4))] tree = flatten_util.ravel_pytree(arr)[0] print(tree) # > [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]

anh-tong · 2023-07-13T05:45:32Z

src/signax/signatures.py

+        return [
+            sig_fun(path[i]) for i in range(path.shape[0])
+        ]  # otherwise, use scan to handle the batch dimension in the list


I think this is not necessary after we handle flatten correctly in the _signature function.

so i did this because i couldn't figure out how to get the batch dimension as the dimension of the list! It seems like vmap only correctly broadcasts the batch dimension to the arrays, but maybe I fed things in wrong. This is the only thing I did that gave the correct shape for now.

this problem does not exist when flatten=True, since it's arrays all the way down!

anh-tong · 2023-07-13T05:57:06Z

src/signax/signatures.py

 def logsignature(
-    path: Float[Array, "path_len dim"], depth: int, stream: bool = False
+    path: Float[Array, "path_len dim"],


I guess you want to make logsignature to handle the input (batch_size, path_len, dim) as well.

In this case, we also need to handle the shape and need an extra vmap

this should work automatically by dispatching to those checks in signature, which should do the right things I think! But let me know if there's logic I'm missing here.

anh-tong · 2023-07-13T05:58:16Z

tests/test_against_signatory.py

+    path = rng.standard_normal(length)
+    signature(path, depth)


We can remove this. This actually causes error since path is 1D.

anh-tong · 2023-07-13T06:07:56Z

Overall, it looks good. When the behaviour of flatten in _signature function is corrected and allowing logsignature to handle inputs with (batch_size, path_len, dim) is added, we can merge this to main.

Also, can you please check why CI format has errors?

phinate · 2023-07-16T17:31:05Z

ok thinks look mostly okay now, i think i've caught the majority of errors -- there are still a couple tests failing with stream=True -- do you think i'm handling the flattening logic wrong somewhere?

anh-tong · 2023-07-17T06:27:31Z

From the comment

signax/src/signax/signatures.py

Line 72 in feb0639

# if not flattening, use list comprehension since vmap cant handle list outputs

I'm curious about why we cannot vmap a function returning a list.

I test with a simple example like this:

import jax
import jax.numpy as jnp
from jax import flatten_util

def fn(x, flatten=False):
    n = x.shape[0]
    res = [jnp.ones((n, )), jnp.ones((n, n))]
    if flatten:
        res = flatten_util.ravel_pytree(res)[0]
    return res

# normal case
x = jnp.ones((2,))
output = fn(x)
[print(o.shape) for o in output]
# returned shape: [(2,), (2, 2)]

# vmap
print("--- vmap ---")
batch_size = 5
x = jnp.ones((batch_size, 2))
output = jax.vmap(fn)(x)
[print(o.shape) for o in output]
# return [(batch_size, 2), (batch_size, 2, 2)]


# flatten + vmap
print("--- flatten + vmap ---")
batch_size = 5
x = jnp.ones((batch_size, 2))
output = jax.vmap(lambda _x: fn(_x, flatten=True))(x)
print(output.shape)
# return (batch_size, 2 + 2 * 2)

We can make _signature handle all the arguments (stream, flatten, ...). The remaining job of signature function is just to check the shape and dispatch between _signature and vmapped _signature

I suggest to handle flatten for the case of stream=True and stream=False in this line:

signax/src/signax/signatures.py

Lines 115 to 123 in feb0639

    
           if stream: 
        
               res = [ 
        
                   jnp.concatenate([first[None, ...], rest], axis=0) 
        
                   for first, rest in zip(exp_term, stacked) 
        
               ] 
        
           else: 
        
               res = carry 
        
           if flatten: 
        
               res = flatten_util.ravel_pytree(res)[0]

with something like this

if stream:
      res = [
          jnp.concatenate([first[None, ...], rest], axis=0)
          for first, rest in zip(exp_term, stacked)
      ]
    # here `res` has shape [(patch_len - 1, dim), (patch_len - 1, dim, dim), ...]
    if flatten:
        res =jax.vmap(lambda _x: flatten_util.ravel_pytree(x)[0])(res)
        # now `res` has shape (patch_len -1, dim + dim * dim + ...)
  else:
      res = carry
      # `res` has shape [(dim,), (dim, dim), ...]
    if flatten:
        res = flatten_util.ravel_pytree(res)[0] # `res` has shape (dim + dim * dim + ..., )

I believe the following line works well. For consistency, probably we can change from signature function to _signature which designates to handle 2D input.

signax/src/signax/signatures.py

Line 214 in feb0639

remainder_signature = signature(path_remainder, depth, stream)

phinate · 2023-07-19T14:10:29Z

I'm curious about why we cannot vmap a function returning a list.
I test with a simple example like this:

import jax
import jax.numpy as jnp
from jax import flatten_util

def fn(x, flatten=False):
    n = x.shape[0]
    res = [jnp.ones((n, )), jnp.ones((n, n))]
    if flatten:
        res = flatten_util.ravel_pytree(res)[0]
    return res

# normal case
x = jnp.ones((2,))
output = fn(x)
[print(o.shape) for o in output]
# returned shape: [(2,), (2, 2)]

# vmap
print("--- vmap ---")
batch_size = 5
x = jnp.ones((batch_size, 2))
output = jax.vmap(fn)(x)
[print(o.shape) for o in output]
# return [(batch_size, 2), (batch_size, 2, 2)]


# flatten + vmap
print("--- flatten + vmap ---")
batch_size = 5
x = jnp.ones((batch_size, 2))
output = jax.vmap(lambda _x: fn(_x, flatten=True))(x)
print(output.shape)
# return (batch_size, 2 + 2 * 2)

It's the second case that I implemented this line for -- would you not expect [(2), (2, 2)]*batch_size instead of [(batch_size, 2), (batch_size, 2, 2)] here? That's the hard thing to do with vmap.

I suggest to handle flatten for the case of stream=True and stream=False in this line:

signax/src/signax/signatures.py

Lines 115 to 123 in feb0639

    
           if stream: 
        
               res = [ 
        
                   jnp.concatenate([first[None, ...], rest], axis=0) 
        
                   for first, rest in zip(exp_term, stacked) 
        
               ] 
        
           else: 
        
               res = carry 
        
           if flatten: 
        
               res = flatten_util.ravel_pytree(res)[0]

with something like this

if stream:
      res = [
          jnp.concatenate([first[None, ...], rest], axis=0)
          for first, rest in zip(exp_term, stacked)
      ]
    # here `res` has shape [(patch_len - 1, dim), (patch_len - 1, dim, dim), ...]
    if flatten:
        res =jax.vmap(lambda _x: flatten_util.ravel_pytree(x)[0])(res)
        # now `res` has shape (patch_len -1, dim + dim * dim + ...)
  else:
      res = carry
      # `res` has shape [(dim,), (dim, dim), ...]
    if flatten:
        res = flatten_util.ravel_pytree(res)[0] # `res` has shape (dim + dim * dim + ..., )

cool, I'll look at this now -- thanks for your help @anh-tong :D

codecov · 2023-07-21T10:33:58Z

Codecov Report

Merging #35 (dd64894) into main (b867de3) will increase coverage by 1.08%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main      #35      +/-   ##
==========================================
+ Coverage   92.85%   93.93%   +1.08%     
==========================================
  Files           7        7              
  Lines         224      264      +40     
==========================================
+ Hits          208      248      +40     
  Misses         16       16

Files Changed	Coverage Δ
src/signax/__init__.py	`100.00% <ø> (ø)`
src/signax/utils.py	`75.92% <ø> (-1.67%)`	⬇️
src/signax/module.py	`100.00% <100.00%> (ø)`
src/signax/signatures.py	`100.00% <100.00%> (ø)`

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

phinate · 2023-07-21T10:37:45Z

stream fix worked! but the code's a bit messy now -- would need to think about how to handle things for logsignature (i currently force no flattening since signature_to_logsignature needs things in list form), but I could probably re-write the flatten logic to happen once in the same way for all functions as opposed to pasting the same code everywhere...

anh-tong · 2023-07-25T04:37:04Z

Hi @phinate, sorry for taking so long. I have updated the code in the following:

function _signature, _logsignature just handle the input with shape (length, channels)
function signature and logsignature built on top of _signature and _logsignature just perform shape handling.
flatten=True now is the default option like signatory

Indeed, keeping track of the shape can be challenging, especially when considering the stream and flatten options being turned on or off.

Let me know if the behavior of current version is what you expect.

anh-tong · 2023-07-28T03:09:38Z

Let's merge this. Thanks for PR.

phinate and others added 4 commits July 11, 2023 15:33

one signature function to rule them all

548b581

another update

5930a10

works

dfdd0e2

confusing commit, revert me if needed

bc58140

anh-tong reviewed Jul 13, 2023

View reviewed changes

phinate and others added 9 commits July 14, 2023 12:23

updates

1c216bb

fix issues

d29890b

type annotation on logsig

b9f7578

try fix tests

5798039

add arg fixes

a89c621

argnames

9195dc8

add flatten to tests against signatory

ffa4f07

shape logic in logsig

e94a64d

flatten correctly in chunked sig

feb0639

phinate added 3 commits July 19, 2023 15:17

anh fix for stream behaviour

ac107ae

fix, but robably not good to always force logsig to not flatten

27925e4

change res to bulk_signature in chunked signature

2f3de45

anh-tong added 2 commits July 25, 2023 13:12

fix shape handling

4f1754d

fix example as api change

88b2c31

improve codecov

c0c99ec

improve codecov

dd64894

anh-tong approved these changes Jul 28, 2023

View reviewed changes

anh-tong merged commit 02e8591 into main Jul 28, 2023
7 checks passed

phinate deleted the shape-handling branch August 13, 2023 12:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Move signature implementation to a single function, and catch shape errors #35

Move signature implementation to a single function, and catch shape errors #35

phinate commented Jul 12, 2023

anh-tong Jul 13, 2023

phinate Jul 14, 2023

anh-tong Jul 13, 2023

phinate Jul 14, 2023

phinate Jul 14, 2023

anh-tong Jul 13, 2023

phinate Jul 14, 2023

anh-tong Jul 13, 2023

phinate Jul 14, 2023

anh-tong commented Jul 13, 2023

phinate commented Jul 16, 2023

anh-tong commented Jul 17, 2023 •

edited

Loading

phinate commented Jul 19, 2023

codecov bot commented Jul 21, 2023 •

edited

Loading

phinate commented Jul 21, 2023

anh-tong commented Jul 25, 2023 •

edited

Loading

anh-tong commented Jul 28, 2023

Move signature implementation to a single function, and catch shape errors #35

Move signature implementation to a single function, and catch shape errors #35

Conversation

phinate commented Jul 12, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

anh-tong commented Jul 13, 2023

phinate commented Jul 16, 2023

anh-tong commented Jul 17, 2023 • edited Loading

phinate commented Jul 19, 2023

codecov bot commented Jul 21, 2023 • edited Loading

Codecov Report

phinate commented Jul 21, 2023

anh-tong commented Jul 25, 2023 • edited Loading

anh-tong commented Jul 28, 2023

anh-tong commented Jul 17, 2023 •

edited

Loading

codecov bot commented Jul 21, 2023 •

edited

Loading

anh-tong commented Jul 25, 2023 •

edited

Loading