The issue here is the mismatch in the standard deviation of train_images vs the first hidden layer. The

train_images.std()
>> 90.0
model.layers[1].get_weights()[0].std()
>> 0.05

0.05 comes from the default std of the kernel initializer of the Dense layer. You need to move these two numbers closer. Ideally in the same order of magnitude, around or smaller than 1.0.

255 isn't a magic number. You can use 100 or 1000 too. You can also specify the desired std of the kernel initializer so that it matches std of your data.

E.g. The following will achieve the desired high accuracy in model.fit() just as well.

train_images = train_images / 100.0  # std = 1.0
train_images.std()
>> 0.9

from tensorflow.keras import initializers
ki = initializers.RandomNormal(stddev=1.0)

model = keras.models.Sequential([
  keras.layers.Flatten(input_shape=[28,28]),
  keras.layers.Dense(300, activation="relu", kernel_initializer=ki), 
  keras.layers.Dense(100, activation="relu"),
  keras.layers.Dense(10, activation="softmax")
])

The issue lies in the mismatch between the standard deviations of train_images vs that of the first hidden layer. The 0.05 comes from the default std of the kernel initializer of the Dense layer. The goal is to get these values closer to each other, Ideally in the same order of magnitude (around or smaller than 1.0).

train_images.std()
>> 90.0
model.layers[1].get_weights()[0].std()
>> 0.05

255 isn't a magic number. You can use 100 or 1000 too. You can also specify the desired std of the kernel initializer so that it matches std of your data.

E.g. The following will achieve the desired high accuracy in model.fit() just as well.

train_images = train_images / 100.0  # std = 1.0
train_images.std()
>> 0.9

from tensorflow.keras import initializers
ki = initializers.RandomNormal(stddev=1.0)

model = keras.models.Sequential([
  keras.layers.Flatten(input_shape=[28,28]),
  keras.layers.Dense(300, activation="relu", kernel_initializer=ki), 
  keras.layers.Dense(100, activation="relu"),
  keras.layers.Dense(10, activation="softmax")
])