fix: resolve NameError when MultiHeadAttention is called with w_init=None

[exopoiesis]

Problem

MultiHeadAttention.__init__ references w_init_scale in the fallback branch:

if w_init is None:
    w_init = hk.initializers.VarianceScaling(w_init_scale)  # NameError

w_init_scale is never defined in this scope, so any call with the default
w_init=None raises NameError: name 'w_init_scale' is not defined.

Fix

Replace the undefined variable with the literal 1.0, which matches the
upstream haiku VarianceScaling default:

if w_init is None:
    w_init = hk.initializers.VarianceScaling(1.0)

Tests

Two focused regression tests in tests/test_attention.py:

• test_w_init_none_does_not_raise — exercises the formerly-broken default path
• test_w_init_explicit_still_works — confirms explicit w_init is unaffected

Scope

This PR is intentionally minimal: only crystalformer/src/attention.py and
tests/test_attention.py are changed. No tree API or extension module changes.

Fixes #11

[zdcao121]

Hi @exopoiesis, thanks for the PR and the detailed tests.

Two notes from our side:

1. jax.tree_util.tree_map is an alias of jax.tree.map() for backward compatibility with older JAX versions, so this replacement is not strictly required for us right now.
2. The extension module is deprecated, so changes there are lower priority.

We still appreciate the w_init=None fix and the overall maintenance effort.

[exopoiesis]

Slimmed down to just the w_init fix + regression test as suggested. Tree_map and extension-module changes dropped.