From 8ebb8a96611bfe388a3072dabd8bb7ea383121f5 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Wed, 1 Oct 2025 09:21:55 -0700
Subject: [PATCH] finished a first pass at digesting the paper, start with
 transformer

---
 dreamer4/dreamer4.py | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/dreamer4/dreamer4.py b/dreamer4/dreamer4.py
index c152ada..972833b 100644
--- a/dreamer4/dreamer4.py
+++ b/dreamer4/dreamer4.py
@@ -1,8 +1,22 @@
+from __future__ import annotations
+
+import math
+from functools import partial
+
 import torch
 import torch.nn.functional as F
-from torch.nn import Module, ModuleList, RMSNorm, Identity
+from torch.nn import Module, ModuleList, Linear, RMSNorm, Identity
 from torch import cat, stack, tensor, Tensor, is_tensor
 
+# ein related
+
+from einops import einsum, rearrange, repeat, reduce
+from einops.layers.torch import Rearrange
+
+# constants
+
+LinearNoBias = partial(Linear, bias = False)
+
 # helpers
 
 def exists(v):