dataset.ChessMoEFactorizedDataset Class Reference

Inheritance diagram for dataset.ChessMoEFactorizedDataset:

Collaboration diagram for dataset.ChessMoEFactorizedDataset:

Public Member Functions
	__init__ (self, data_dir, max_samples=None, in_memory=False, n_globals=21)
	__len__ (self)
	__getitem__ (self, idx)

Public Attributes
	in_memory = in_memory
	data_dir = data_dir
	n_globals = n_globals
	features_path = os.path.join(data_dir, "features.bin")
	labels_path = os.path.join(data_dir, "labels.bin")
	weights_path = os.path.join(data_dir, "expert_weights.bin")
	features_dtype_legacy
	features_dtype_compact
	features_dtype = self.features_dtype_compact
str	feature_layout = "compact54"
	n_samples = expected_samples
	features = None
	labels = None
	weights = None

Static Public Attributes
	PLANES_PER_TYPE = PLANES_PER_GROUP_CONST
int	MAX_BRANCH_PLANES = 10
	PACKED_BRANCH_PLANES = sum(PLANES_PER_TYPE)
int	NUM_BYPASS_PLANES = 12
	PACKED_OFFSETS = PACKED_OFFSETS_CONST

Protected Member Functions
	_load_to_memory (self)
	_lazy_load (self)

Detailed Description

Dataset for factorized MoE training.

Returns dict with:
- branches: [54, 8, 8] uint8 (compact packed branch planes)
  - bypass: [12, 8, 8] uint8 (packed bypass planes)
  - global: [n_globals] float32 (default first 21 globals)
  - wdl: [3] float32
  - base_weights: [4] float32
  - aux_weights: [2] float32

Constructor & Destructor Documentation

__init__()

dataset.ChessMoEFactorizedDataset.__init__	(		self,
			data_dir,
			max_samples = None,
			in_memory = False,
			n_globals = 21 )

Member Function Documentation

__getitem__()

dataset.ChessMoEFactorizedDataset.__getitem__	(		self,
			idx )

Here is the call graph for this function:

__len__()

dataset.ChessMoEFactorizedDataset.__len__

(

self

)

_lazy_load()

dataset.ChessMoEFactorizedDataset._lazy_load ( self )

protected

Here is the caller graph for this function:

_load_to_memory()

dataset.ChessMoEFactorizedDataset._load_to_memory ( self )

protected

Member Data Documentation

data_dir

dataset.ChessMoEFactorizedDataset.data_dir = data_dir

feature_layout

str dataset.ChessMoEFactorizedDataset.feature_layout = "compact54"

features

dataset.ChessMoEFactorizedDataset.features = None

features_dtype

dataset.ChessMoEFactorizedDataset.features_dtype = self.features_dtype_compact

features_dtype_compact

dataset.ChessMoEFactorizedDataset.features_dtype_compact

Initial value:

=  np.dtype(
            [
                ("branches", np.uint64, self.PACKED_BRANCH_PLANES),
                ("bypass_continuous", np.uint8, (2, 64)),
                ("bypass_categorical", np.uint64, 10),
                ("global", np.float32, 32),
            ],
            align=True,
        )

features_dtype_legacy

dataset.ChessMoEFactorizedDataset.features_dtype_legacy

Initial value:

=  np.dtype(
            [
                ("branches", np.uint8, (12, self.MAX_BRANCH_PLANES, 64)),
                ("bypass", np.uint8, (self.NUM_BYPASS_PLANES, 64)),
                ("global", np.float32, 32),
            ],
            align=True,
        )

features_path

dataset.ChessMoEFactorizedDataset.features_path = os.path.join(data_dir, "features.bin")

in_memory

dataset.ChessMoEFactorizedDataset.in_memory = in_memory

labels

dataset.ChessMoEFactorizedDataset.labels = None

labels_path

dataset.ChessMoEFactorizedDataset.labels_path = os.path.join(data_dir, "labels.bin")

MAX_BRANCH_PLANES

int dataset.ChessMoEFactorizedDataset.MAX_BRANCH_PLANES = 10

static

n_globals

dataset.ChessMoEFactorizedDataset.n_globals = n_globals

n_samples

dataset.ChessMoEFactorizedDataset.n_samples = expected_samples

NUM_BYPASS_PLANES

int dataset.ChessMoEFactorizedDataset.NUM_BYPASS_PLANES = 12

static

PACKED_BRANCH_PLANES

dataset.ChessMoEFactorizedDataset.PACKED_BRANCH_PLANES = sum(PLANES_PER_TYPE)

static

PACKED_OFFSETS

dataset.ChessMoEFactorizedDataset.PACKED_OFFSETS = PACKED_OFFSETS_CONST

static

PLANES_PER_TYPE

dataset.ChessMoEFactorizedDataset.PLANES_PER_TYPE = PLANES_PER_GROUP_CONST

static

weights

dataset.ChessMoEFactorizedDataset.weights = None

weights_path

dataset.ChessMoEFactorizedDataset.weights_path = os.path.join(data_dir, "expert_weights.bin")

---

The documentation for this class was generated from the following file:

• training/dataset.py