Flax는 Jax와 연동되는 library.기본적인 Network architecture를 제공 [e.g. Dense, Conv ,...]간단하게 torch.nn.Linear, torch.nn.Conv1d,.. 를 사용할 수 있도록 함. 대표적인 사용 예시는 아래와 같음1. Dense Layer [torch.nn.Linear]import jaximport flax.linen as nnfrom typing import Sequnece, Callabledef default_init(scale: Optional[float] = jnp.sqrt(2)): return nn.initializers.orthogonal(scale)class Linear(nn.Module) hidden_layers: Sequnec..

import jaxprint(jax.devices())위 코드를 실행하면 아래와 같은 에러가 나왔음.WARNING: Logging before InitGoogle() is written to STDERRI0000 00:00:1721141137.213010 8925 singleton_tpu_states_manager.cc:45] TpuStatesManager::GetOrCreate(): no tpu system exists. Creating a new tpu system.F0000 00:00:1721141137.223216 8925 b295d63588a.cc:1205] File: InitGoogle() has not finished yet. See gl_______________________..