GitHub - PJHkorea/egregore-core-jax: High-performance, branchless numerical stability kernels and compiler-optimized core infrastructure for advanced JAX/XLA deep learning architectures. · GitHub

/" data-turbo-transient="true" />

Skip to content

Search or jump to...

Search code, repositories, users, issues, pull requests...

-->

Search

Clear

Search syntax tips

Provide feedback

--><br>We read every piece of feedback, and take your input very seriously.

Include my email address so I can be contacted

Cancel

Submit feedback

Saved searches

Use saved searches to filter your results more quickly

-->

Name

Query

To see all available qualifiers, see our documentation.

Cancel

Create saved search

Sign in

/;ref_cta:Sign up;ref_loc:header logged out"}"<br>Sign up

Appearance settings

Resetting focus

You signed in with another tab or window. Reload to refresh your session.<br>You signed out in another tab or window. Reload to refresh your session.<br>You switched accounts on another tab or window. Reload to refresh your session.

Dismiss alert

{{ message }}

PJHkorea

egregore-core-jax

Public

Notifications<br>You must be signed in to change notification settings

Fork

Star

main

BranchesTags

Go to file

CodeOpen more actions menu

Folders and files<br>NameNameLast commit message<br>Last commit date<br>Latest commit

History<br>25 Commits<br>25 Commits

egregore_core_jax

egregore_core_jax

LICENSE

LICENSE

README.md

README.md

pyproject.toml

pyproject.toml

View all files

Repository files navigation

🧩 Core Modules Specifications (핵심 모듈 명세)

1. 🛠️ optimizers.py — Branchless PyTree LLRD Engine

JAX/Optax 유저들을 위해 인터프리터의 조건 분기 개입을 제거한 초고속 계층별 차등 최적화(LLRD) 팩토리 모듈입니다.

[KR] 목적 및 기능 : 복잡한 PyTree 가중치 사전 스캔 시 발생하던 파이썬 호스트 오버헤드와 hasattr 분기 예측 실패(Branch Stall)를 'getattr 오리 타이핑 속성 마스킹'과 '삼항 대수식 마스크'를 통해 평탄화했습니다. 빈 경로 유입 크래시를 방어하는 safe_path 융합 기술이 탑재되어 있으며, custom_router 의존성 주입(DI)을 지원하여 LLM 등 모든 JAX 모델에 범용 결합을 위함이 목적입니다.

[EN] Purpose & Function : A high-velocity Layer-wise Learning Rate Decay (LLRD) factory module crafted to eliminate conversational python interpreter overhead for JAX/Optax developers. It flattens the conventional hasattr branch stalls during complex PyTree scanning by utilizing 'getattr Duck-Typing Attribute Masking' and inline algebraic string-matching. Equipped with a robust safe_path concatenation technique to safeguard against empty path crashes, it offers a custom_router dependency injection interface for universal integration into any JAX-based architecture, including LLMs.

💡 TL;DR (핵심 요약) :

[KR] 문제와 해법 : JAX에서 레이어별로 학습률을 다르게 줄 때(LLRD), 가중치 트리 경로를 스캔하느라 파이썬의 if/else 조건문을 쓰게 되면 GPU/TPU 컴파일 그래프가 조각나고 병목이 발생합니다. 이 모듈은 조건 분기를 완전히 파괴하고 "gate" * (is_gate) + "backbone" * (not_is_gate) 같은 인라인 대수식 문자열 연산 만으로 파라미터를 하드웨어 레벨에서 초고속 라우팅합니다.

[EN] Problem & Solution : When implementing Layer-wise Learning Rate Decay (LLRD) in JAX, using standard Python if/else predicates to scan weight tree paths fragments the GPU/TPU compilation graph and triggers severe branch stalls. This module thoroughly obliterates conditional branching, utilizing inline algebraic string-manipulation like "gate" * (is_gate) + "backbone" * (not_is_gate) to execute ultra-high-speed parameter routing natively within hardware-level constraints.

🚀 [Update] Production-Ready Enterprise Silicon MUX Architecture :

[KR] 진화된 핵심 사양 : 상위 프레임워크(Optax)의 문자열 기반 딕셔너리 라우팅 제약마저 전산학적으로 완전히 박멸했습니다. 불린 결과값(True/False)에 곧바로 * 1.0f 대수 연산을 유도하여 가속기 ALU 단축 레지스터가 소모할 수 있는 float32 정적 리터럴 마스크로 즉시 압축(is_gate = (root_key_name == "gate") * 1.0f)합니다. 다중 옵티마이저 분기 체인을 완전히 소멸시키고, 단일 기저 AdamW 위에서 가속기 내부의 단일 클럭 아다마르 곱(Hadamard Product) 텐서 연산만으로 업데이트 스케일을 제어하여 XLA 컴파일러가 온칩 SRAM 내부에 최적의 단일 융합 커널(Fused Kernel)을 강제 포메이션하도록 리엔지니어링되었습니다.

[EN] Advanced Evolution Mechanics : Fundamentally extinguishes framework-level string dictionary routing. By directly appending * 1.0f to boolean evaluations, it compresses dynamic states into float32 literal register masks (is_gate = (root_key_name == "gate") * 1.0f) tailored for the accelerator ALU. This demolishes multi-optimizer branch chains, consolidating execution onto a unified underlying AdamW engine where layer-wise scaling is driven strictly via single-cycle inline Hadamard tensor products, forcing the XLA compiler to generate high-density Fused Kernels directly inside on-chip SRAM.

2. 📐 geometry.py — Static Virtual View Curvature Sensor

다차원 배치 입력의 유동적인 확장 속에서도 정적 컴파일 무결성을 목표로 하는 고속 다양체 곡률(공간 분산) 계측 모듈입니다.

[KR] 목적 및 기능 : JAX 환경에서 jnp.reshape나 차원 튜플 조립 시 유저들의 골치를 썩이던 형상 추적 에러(ConcretizationTypeError)를 static_argnums=(1,) 정적 상수 락(Lock) 가드레일로 막아봤습니다. 임의의 N차원 배치 입력([B, D], [B, T, D] 등)이 인입되더라도 단일 정적 가상 뷰 레벨의 인라인 연속 메모리 레이아웃으로 변환하여 하드웨어 온칩 리덕션 속도를 극대화함이 목적입니다.

[EN] Purpose & Function : A high-efficiency spatial macro-curvature (geometric variance proxy) measurement module engineered to maintain static compilation integrity across fluid multi-dimensional batch scaling. To shield developers from stubborn tracer crashes...