Added dtype arg to jax.random.normal, refactored blend_v, blend_h, blend_t methods, replaced 1/jnp.sqrt with jax.lax.rsqrt

Author: prishajain1

src/maxdiffusion/models/ltx2/autoencoder_kl_ltx2.py

@@ -51,8 +51,7 @@ def __call__(self, x: jax.Array, channel_dim: Optional[int] = None) -> jax.Array
     channel_dim = channel_dim if channel_dim is not None else self.channel_dim
     # Compute mean of squared values along channel dimension.
     mean_sq = jnp.mean(jnp.square(x), axis=channel_dim, keepdims=True)
-    rms = jnp.sqrt(mean_sq + self.eps)
-    return x / rms
+    return x * jax.lax.rsqrt(mean_sq + self.eps)
 
 
 class LTX2VideoCausalConv3d(nnx.Module):
@@ -226,7 +225,7 @@ def __init__(
       self.per_channel_scale2 = None
 
     if timestep_conditioning:
-      self.scale_shift_table = nnx.Param(jax.random.normal(rngs.params(), (4, in_channels)) / (in_channels**0.5))
+      self.scale_shift_table = nnx.Param(jax.random.normal(rngs.params(), (4, in_channels), dtype=dtype) / (in_channels**0.5))
     else:
       self.scale_shift_table = None
 
@@ -1261,24 +1260,42 @@ def enable_tiling(
 
   def blend_v(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[2], b.shape[2], blend_extent)
-    for y in range(blend_extent):
-      val = a[:, :, -blend_extent + y, :, :] * (1 - y / blend_extent) + b[:, :, y, :, :] * (y / blend_extent)
-      b = b.at[:, :, y, :, :].set(val)
-    return b
+    if blend_extent <= 0:
+        return b
+        
+    # Create broadcastable blending weights: (1, 1, blend_extent, 1, 1)
+    y = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, 1, -1, 1, 1)
+    
+    val = a[:, :, -blend_extent:, :, :] * (1.0 - y / blend_extent) + \
+          b[:, :, :blend_extent, :, :] * (y / blend_extent)
+          
+    return b.at[:, :, :blend_extent, :, :].set(val)
 
   def blend_h(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[3], b.shape[3], blend_extent)
-    for x in range(blend_extent):
-      val = a[:, :, :, -blend_extent + x, :] * (1 - x / blend_extent) + b[:, :, :, x, :] * (x / blend_extent)
-      b = b.at[:, :, :, x, :].set(val)
-    return b
+    if blend_extent <= 0:
+        return b
+        
+    # Create broadcastable blending weights: (1, 1, 1, blend_extent, 1)
+    x = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, 1, 1, -1, 1)
+    
+    val = a[:, :, :, -blend_extent:, :] * (1.0 - x / blend_extent) + \
+          b[:, :, :, :blend_extent, :] * (x / blend_extent)
+          
+    return b.at[:, :, :, :blend_extent, :].set(val)
 
   def blend_t(self, a: jax.Array, b: jax.Array, blend_extent: int) -> jax.Array:
     blend_extent = min(a.shape[1], b.shape[1], blend_extent)
-    for x in range(blend_extent):
-      val = a[:, -blend_extent + x, :, :, :] * (1 - x / blend_extent) + b[:, x, :, :, :] * (x / blend_extent)
-      b = b.at[:, x, :, :, :].set(val)
-    return b
+    if blend_extent <= 0:
+        return b
+        
+    # Create broadcastable blending weights: (1, blend_extent, 1, 1, 1)
+    x = jnp.arange(blend_extent, dtype=a.dtype).reshape(1, -1, 1, 1, 1)
+    
+    val = a[:, -blend_extent:, :, :, :] * (1.0 - x / blend_extent) + \
+          b[:, :blend_extent, :, :, :] * (x / blend_extent)
+          
+    return b.at[:, :blend_extent, :, :, :].set(val)
 
   def tiled_encode(self, x: jax.Array, key: Optional[jax.Array] = None, causal: Optional[bool] = None) -> jax.Array:
     B, T, H, W, C = x.shape