Pass max_length_time_axis instead of max_size

Makes it so that the warning:

```
Setting max_size dynamically sets the `max_length_time_axis` to be `max_size`//`add_batch_size = .*`
```

will no longer be triggered by legitimate use of `create_flat_buffer` and
`make_prioritised_flat_buffer`.

flashbax/buffers/flat_buffer.py

@@ -123,13 +123,13 @@ def create_flat_buffer(
         )
 
         buffer = make_trajectory_buffer(
-            max_length_time_axis=None,  # Unused because max_size is specified
+            max_length_time_axis=max_length // add_batch_size,
             min_length_time_axis=min_length // add_batch_size + 1,
             add_batch_size=add_batch_size,
             sample_batch_size=sample_batch_size,
             sample_sequence_length=2,
             period=1,
-            max_size=max_length,
+            max_size=None,
         )
 
     add_fn = buffer.add

flashbax/buffers/prioritised_flat_buffer.py

@@ -110,13 +110,13 @@ def make_prioritised_flat_buffer(
         )
 
         buffer = make_prioritised_trajectory_buffer(
-            max_length_time_axis=None,  # Unused because max_size is specified
+            max_length_time_axis=max_length // add_batch_size,
             min_length_time_axis=min_length // add_batch_size + 1,
             add_batch_size=add_batch_size,
             sample_batch_size=sample_batch_size,
             sample_sequence_length=2,
             period=1,
-            max_size=max_length,
+            max_size=None,
             priority_exponent=priority_exponent,
             device=device,
         )

pyproject.toml

@@ -16,7 +16,6 @@ filterwarnings = [
     "error",
     "ignore:`sample_sequence_length` greater than `min_length_time_axis`:UserWarning:flashbax",
     "ignore:Setting period greater than sample_sequence_length will result in no overlap betweentrajectories:UserWarning:flashbax",
-    "ignore:Setting max_size dynamically sets the `max_length_time_axis` to be `max_size`//`add_batch_size = .*`:UserWarning:flashbax",
     "ignore:jax.tree_map is deprecated:DeprecationWarning:flashbax",
 ]