class_docstr: A wrapper around torchsde.BrownianTree that enables batches of entropy.
class_docstr_tok: ["A","wrapper","around","torchsde.BrownianTree","that","enables","batches","of","entropy","."]
code: def __init__(self, x, t0, t1, seed=None, **kwargs):
      t0, t1, self.sign = self.sort(t0, t1)
      w0 = kwargs.get("w0", torch.zeros_like(x))
      if seed is None:
          seed = torch.randint(0, 2**63 - 1, []).item()
      self.batched = True
      try:
          assert len(seed) == x.shape[0]
          w0 = w0[0]
      except TypeError:
          seed = [seed]
          self.batched = False
      self.trees = [
          torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs)
          for s in seed
      ]
  
imports: import math
  from collections import defaultdict
  from typing import List, Optional, Tuple, Union
  import numpy
  import torch
  import torchsde
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
label_desc: Use image node_id 1 to create a new BatchedBrownianTree object  with example: obj = BatchedBrownianTree(x, t0, t1, seed)