Enable pytree step_size, offsets input

[ASEM000]

import finitediffx as fdx
import jax.numpy as jnp

params = {"a": 1.0, "b": 2.0, "c": 3.0}


def f(params):
    return params["a"] ** 2 + params["b"]


step_size = {"a": 1, "b": 1, "c": 1}
offsets = {
    "a": jnp.array([-1, 1]),
    "b": jnp.array([-1, 1]),
    "c": jnp.array([-1, 1]),
}

fdx.fgrad(f, step_size=step_size, offsets=offsets)(params)
# {'a': Array(2., dtype=float32),
#  'b': Array(1., dtype=float32),
#  'c': Array(0., dtype=float32)}

[codecov]

Codecov Report

Patch coverage: 93.02% and project coverage change: -3.31 ⚠️

Comparison is base (a648314) 98.39% compared to head (33a7dff) 95.08%.

❗ Current head 33a7dff differs from pull request most recent head e145aa8. Consider uploading reports for the commit e145aa8 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main       #5      +/-   ##
==========================================
- Coverage   98.39%   95.08%   -3.31%     
==========================================
  Files           6        6              
  Lines         624      672      +48     
==========================================
+ Hits          614      639      +25     
- Misses         10       33      +23     

Impacted Files	Coverage Δ
finitediffx/_src/fgrad.py	80.12% <88.88%> (-13.36%)	⬇️
finitediffx/_src/utils.py	97.87% <100.00%> (ø)
tests/test_fgrad.py	100.00% <100.00%> (ø)

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.